Update huggingface_utils.py

ServiceNow · Nov 29, 2024 · 5922437 · 5922437
1 parent c16a69d
commit 5922437
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/agentlab/llm/huggingface_utils.py b/src/agentlab/llm/huggingface_utils.py
@@ -6,6 +6,7 @@
 from transformers import AutoTokenizer, GPT2TokenizerFast
 
 from agentlab.llm.base_api import AbstractChatModel
+from agentlab.llm.llm_utils import Discussion
 from agentlab.llm.prompt_templates import PromptTemplate, get_prompt_template
 
 
@@ -59,7 +60,8 @@ def __call__(
         if self.tokenizer:
             # messages_formated = _convert_messages_to_dict(messages) ## ?
             try:
-                messages.merge()
+                if isinstance(messages, Discussion):
+                    messages.merge()
                 prompt = self.tokenizer.apply_chat_template(messages, tokenize=False)
             except Exception as e:
                 if "Conversation roles must alternate" in str(e):