debugging log_probs for hf models

Author: TLSDC

src/agentlab/llm/chat_api.py

@@ -145,6 +145,7 @@ def make_model(self):
                 temperature=self.temperature,
                 max_new_tokens=self.max_new_tokens,
                 n_retry_server=self.n_retry_server,
+                log_probs=self.log_probs
             )
         else:
             raise ValueError(f"Backend {self.backend} is not supported")
@@ -237,7 +238,7 @@ def __init__(
         self.max_tokens = max_tokens
         self.max_retry = max_retry
         self.min_retry_wait_time = min_retry_wait_time
-        self.logprobs = log_probs
+        self.log_probs = log_probs
 
         # Get the API key from the environment variable if not provided
         if api_key_env_var:
@@ -284,7 +285,7 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
                     n=n_samples,
                     temperature=temperature,
                     max_tokens=self.max_tokens,
-                    logprobs=self.logprobs,
+                    log_probs=self.log_probs,
                 )
 
                 if completion.usage is None:
@@ -315,8 +316,8 @@ def __call__(self, messages: list[dict], n_samples: int = 1, temperature: float
 
         if n_samples == 1:
             res = AIMessage(completion.choices[0].message.content)
-            if self.logprobs:
-                res["logprobs"] = completion.choices[0].logprobs
+            if self.log_probs:
+                res["log_probs"] = completion.choices[0].log_probs
             return res
         else:
             return [AIMessage(c.message.content) for c in completion.choices]
@@ -429,7 +430,7 @@ def __init__(
         n_retry_server: Optional[int] = 4,
         log_probs: Optional[bool] = False,
     ):
-        super().__init__(model_name, base_model_name, n_retry_server)
+        super().__init__(model_name, base_model_name, n_retry_server, log_probs)
         if temperature < 1e-3:
             logging.warning("Models might behave weirdly when temperature is too low.")
         self.temperature = temperature

src/agentlab/llm/huggingface_utils.py

@@ -2,12 +2,11 @@
 import time
 from typing import Any, List, Optional, Union
 
-from pydantic import Field
-from transformers import AutoTokenizer, GPT2TokenizerFast
-
 from agentlab.llm.base_api import AbstractChatModel
 from agentlab.llm.llm_utils import AIMessage, Discussion
 from agentlab.llm.prompt_templates import PromptTemplate, get_prompt_template
+from pydantic import Field
+from transformers import AutoTokenizer, GPT2TokenizerFast
 
 
 class HFBaseChatModel(AbstractChatModel):
@@ -40,9 +39,10 @@ class HFBaseChatModel(AbstractChatModel):
         description="The number of times to retry the server if it fails to respond",
     )
 
-    def __init__(self, model_name, base_model_name, n_retry_server):
+    def __init__(self, model_name, base_model_name, n_retry_server, log_probs):
         super().__init__()
         self.n_retry_server = n_retry_server
+        self.log_probs = log_probs
 
         if base_model_name is None:
             self.tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -102,8 +102,9 @@ def __call__(
                     temperature = temperature if temperature is not None else self.temperature
                     answer = self.llm(prompt, temperature=temperature)
                     response = AIMessage(answer)
-                    if hasattr(answer, "details"):
-                        response["log_prob"] = answer.details.log_prob
+                    if self.log_probs:
+                        response["content"] = answer.generated_text
+                        response["log_prob"] = answer.details
                     responses.append(response)
                     break
                 except Exception as e: