Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions src/agentlab/llm/huggingface_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,12 @@
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):
Expand Down Expand Up @@ -104,7 +105,7 @@ def __call__(
response = AIMessage(answer)
if self.log_probs:
response["content"] = answer.generated_text
response["log_prob"] = answer.details
response["log_probs"] = answer.details
responses.append(response)
break
except Exception as e:
Expand Down