Add logprobs to gpt 3.5

.gitignore

@@ -1,4 +1,3 @@
 build/
 dist/
 *.spec
-uv.lock

README.md

@@ -26,7 +26,7 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
 
 Then, on both your local machine and Sherlock (or either one, if you just want some features):
 ```
-uv tool install "rl[sherlock] @ git+https://github.com/ProbablyFaiz/rl.git@v0.8.1"
+uv tool install "rl[sherlock] @ git+https://github.com/ProbablyFaiz/rl.git@v0.9.0"
 ```
 
 ## Setup

pyproject.toml

@@ -5,7 +5,7 @@ authors = [
     { name = "Faiz Surani", email = "[email protected]" },
     { name = "Varun Magesh", email="[email protected]" }
 ]
-description = "A CLI for various RegLab / Sherlock tasks"
+description = "Utilities for various RegLab / Sherlock tasks"
 readme = "README.md"
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -25,8 +25,8 @@ dependencies = [
     "python-dotenv",
     "unidecode",
     "StrEnum",
-    "watchdog",
     # TODO: Maybe move into an extra?
+    "watchdog",
     "pydantic>=2",
     "typing-extensions",
     "legal-segmenter @ git+https://github.com/lexeme-dev/legal-segmenter@main",
@@ -42,9 +42,9 @@ sherlock = [
     "paramiko",
 ]
 local_llm = [
+    "rl[llm]",
     "transformers>=4.43",
     "torch",
-    "rl[llm]",
     "vllm~=0.5.3post1; platform_system == 'Linux'",
     "peft",
     "trl",

rl/llm/engines/client.py

@@ -44,6 +44,9 @@ def batch_generate(self, prompts: list[ChatInput]) -> list[InferenceOutput]:
         )
 
 
+_NUM_LOGPROBS = 20
+
+
 class _OAIClientEngine(ClientEngine, ABC):
     BASE_URL: str
     API_KEY_NAME: str
@@ -82,16 +85,26 @@ def generate(self, prompt: ChatInput) -> InferenceOutput:
             completion_kwargs["max_tokens"] = self.llm_config.max_new_tokens
         if EngineFeature.JSON_OUTPUT in self.enabled_features:
             completion_kwargs["response_format"] = {"type": "json_object"}
+        if EngineFeature.RETURN_LOGPROBS in self.enabled_features:
+            completion_kwargs["logprobs"] = True
+            completion_kwargs["top_logprobs"] = _NUM_LOGPROBS
 
         response = self.client.chat.completions.create(**completion_kwargs)
-        return InferenceOutput(
+        choice = response.choices[0]
+
+        output = InferenceOutput(
             prompt=prompt,  # type: ignore
-            text=response.choices[0].message.content,
+            text=choice.message.content,
             metadata={
                 "model": self.llm_config.model_name_or_path,
                 "base_url": self.BASE_URL,
             },
         )
+        if EngineFeature.RETURN_LOGPROBS in self.enabled_features:
+            output.logprobs = [
+                {tkn.token: tkn.logprob for tkn in pos.top_logprobs}
+                for pos in choice.logprobs.content
+            ]
 
 
 @register_engine(
@@ -107,7 +120,7 @@ class TogetherEngine(_OAIClientEngine):
 @register_engine(
     "openai",
     required_modules=("openai",),
-    supported_features=(EngineFeature.JSON_OUTPUT,),
+    supported_features=(EngineFeature.JSON_OUTPUT, EngineFeature.RETURN_LOGPROBS),
 )
 class OpenAIEngine(_OAIClientEngine):
     BASE_URL = "https://api.openai.com/v1"
@@ -322,7 +335,7 @@ def _get_modal_app_name(self, model_name: str) -> str:
 @register_engine(
     "batch_openai",
     required_modules=("openai",),
-    supported_features=(EngineFeature.JSON_OUTPUT,),
+    supported_features=(EngineFeature.JSON_OUTPUT, EngineFeature.RETURN_LOGPROBS),
 )
 class BatchOpenAIEngine(InferenceEngine):
     def __init__(self, llm_config: LLMConfig):
@@ -343,6 +356,8 @@ def generate(self, prompt: ChatInput) -> InferenceOutput:
         return self.batch_generate([prompt])[0]
 
     def batch_generate(self, prompts: list[ChatInput]) -> list[InferenceOutput]:
+        from openai.types.chat import ChatCompletion
+
         with tempfile.NamedTemporaryFile(
             mode="w+", suffix=".jsonl", delete=False
         ) as temp_file:
@@ -357,6 +372,9 @@ def batch_generate(self, prompts: list[ChatInput]) -> list[InferenceOutput]:
                     body_kwargs["temperature"] = self.llm_config.temperature
                 if EngineFeature.JSON_OUTPUT in self.enabled_features:
                     body_kwargs["response_format"] = {"type": "json_object"}
+                if EngineFeature.RETURN_LOGPROBS in self.enabled_features:
+                    body_kwargs["logprobs"] = True
+                    body_kwargs["top_logprobs"] = _NUM_LOGPROBS
                 request = {
                     "custom_id": f"request-{i}",
                     "method": "POST",
@@ -396,17 +414,22 @@ def batch_generate(self, prompts: list[ChatInput]) -> list[InferenceOutput]:
 
         outputs = []
         for result in results:
-            response = result["response"]["body"]
-            outputs.append(
-                InferenceOutput(
-                    prompt=prompts[int(result["custom_id"].split("-")[1])],
-                    text=response["choices"][0]["message"]["content"],
-                    metadata={
-                        "model": self.llm_config.model_name_or_path,
-                        "base_url": "https://api.openai.com/v1",
-                    },
-                )
+            parsed_result = ChatCompletion.model_validate(result["response"]["body"])
+            choice = parsed_result.choices[0]
+            output = InferenceOutput(
+                prompt=prompts[int(result["custom_id"].split("-")[1])],
+                text=choice.message.content,
+                metadata={
+                    "model": self.llm_config.model_name_or_path,
+                    "base_url": "https://api.openai.com/v1",
+                },
             )
+            if EngineFeature.RETURN_LOGPROBS in self.enabled_features:
+                output.logprobs = [
+                    {tkn.token: tkn.logprob for tkn in pos.top_logprobs}
+                    for pos in choice.logprobs.content
+                ]
+            outputs.append(output)
 
         Path(temp_file_path).unlink()
 

rl/llm/engines/core.py

@@ -29,7 +29,10 @@ class InferenceOutput(BaseModel):
     prompt: InferenceInput
     text: str
 
-    logprobs: list[dict[int, float]] | None = None
+    # TODO: OpenAI gives us logprobs in terms of the string representation of the token
+    # but vLLM gives us logprobs in terms of the token itself. We should probably
+    # standardize on one representation.
+    logprobs: list[dict[int | str, float]] | None = None
     metadata: dict[str, Any] = Field(default_factory=dict)