Feature/hf pipeline format (#47)

* removed older version utils

* changed requirements

* added changes for inference

.gitignore

@@ -1,4 +1,5 @@
 __pycache__/
 .cache/
 openai_api.key
+hf_token.key
 core

app.py

@@ -12,7 +12,7 @@ with open("openai_api.key", "r") as f:
     os.environ["OPENAI_API_KEY"] = f.read().strip()
 
 DEPLOYED = os.getenv("DEPLOYED", "true").lower() == "true"
-DEFAULT_MODEL_SELECTION = "gpt-3.5-turbo" # "mistralai/Mistral-7B-Instruct-v0.1"
+DEFAULT_MODEL_SELECTION = "gpt-3.5-turbo"
 TEMPERATURE = 0.7
 TOP_P = 1
 MAX_TOKENS = 1024
@@ -147,7 +147,7 @@ def sotopia_info_accordion(accordion_visible=True):
                 interactive=True,
             )
             model_name_dropdown = gr.Dropdown(
-                choices=["cmu-lti/sotopia-pi-mistral-7b-BC_SR", "mistralai/Mistral-7B-Instruct-v0.1", "gpt-3.5-turbo", "gpt-4-turbo"],
+                choices=["cmu-lti/sotopia-pi-mistral-7b-BC_SR", "cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit", "mistralai/Mistral-7B-Instruct-v0.1", "gpt-3.5-turbo", "gpt-4-turbo"],
                 value=DEFAULT_MODEL_SELECTION,
                 interactive=True,
                 label="Model Selection"

langchain_callback_handler.py

@@ -0,0 +1,60 @@
+import logging
+from typing import Any
+
+from langchain.callbacks import StdOutCallbackHandler
+
+logging.addLevelName(15, "LangChain")
+
+
+class LoggingCallbackHandler(StdOutCallbackHandler):
+    """Callback Handler that prints to std out."""
+
+    always_verbose = True
+
+    def __init__(self, name: str) -> None:
+        """Initialize callback handler."""
+        super().__init__()
+        self.logger = logging.getLogger(name)
+        self.prompt = ""
+
+    def on_chain_start(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def on_chain_end(self, *args: Any, **kwargs: Any) -> None:
+        pass
+
+    def on_agent_action(self, *args: Any, **kwargs: Any) -> Any:
+        pass
+
+    def on_tool_end(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        pass
+
+    def on_tool_error(
+        self, error: BaseException | KeyboardInterrupt, **kwargs: Any
+    ) -> None:
+        """Do nothing."""
+        pass
+
+    def on_text(
+        self,
+        text: str,
+        color: str | None = None,
+        end: str = "",
+        **kwargs: Any,
+    ) -> None:
+        """Run when agent ends."""
+        # leave only prompt for environment
+        text = text.replace("\x1b[32;1m\x1b[1;3mHuman: ", "")
+        logging.log(15, f"LLM Call: {text}")
+        self.prompt = text
+
+    def retrive_prompt(self) -> str:
+        return self.prompt
+
+    def on_agent_finish(self, *args: Any, **kwargs: Any) -> None:
+        """Run on agent end."""
+        pass

message_classes.py

@@ -0,0 +1,343 @@
+import re
+from typing import Literal, cast
+
+from pydantic import BaseModel, Field
+
+from utils import format_docstring
+
+ActionType = Literal["none", "speak", "non-verbal communication", "action", "leave"]
+
+
+class Message(BaseModel):
+    """
+    An interface for messages.
+    There is only one required method: to_natural_language
+    """
+
+    def to_natural_language(self) -> str:
+        raise NotImplementedError
+
+
+class SimpleMessage(Message):
+    """
+    A simple message with a single string field.
+    """
+
+    message: str = Field(description="the message")
+
+    def to_natural_language(self) -> str:
+        return self.message
+
+
+class Observation(Message):
+    last_turn: str = Field(description="the last turn of the conversation")
+    turn_number: int = Field(description="the turn number of the conversation")
+    available_actions: list[ActionType] = Field(description="the available actions")
+
+    def to_natural_language(self) -> str:
+        if self.turn_number == 0:
+            return f"\n{self.last_turn}\nConversation Starts:\n"
+        else:
+            return f"Turn #{self.turn_number-1}: {self.last_turn}\n"
+
+
+class ScriptBackground(Message):
+    scenario: str = Field(description="scenario of the episode")
+    p1_name: str = Field(description="name of participant 1")
+    p2_name: str = Field(description="name of participant 2")
+    p1_background: str = Field(description="background of participant 1")
+    p2_background: str = Field(description="background of participant 2")
+    p1_goal: str = Field(description="goal of participant 1")
+    p2_goal: str = Field(description="goal of participant 2")
+
+    def to_natural_language(self) -> str:
+        if self.p1_background or self.p2_background:
+            p1_background = self.p1_background if self.p1_background else "Unknown"
+            p2_background = self.p2_background if self.p2_background else "Unknown"
+            # Not using AND, since in stranger relation the background is not visible
+            return format_docstring(
+                f"""Here is the context of this interaction:
+            Scenario: {self.scenario}
+            Participants: {self.p1_name} and {self.p2_name}
+            {self.p1_name}'s background: {p1_background}
+            {self.p2_name}'s background: {p2_background}
+            {self.p1_name}'s goal: {self.p1_goal}
+            {self.p2_name}'s goal: {self.p2_goal}
+            """
+            )
+        else:
+            return format_docstring(
+                f"""Here is the context of this interaction:
+            Scenario: {self.scenario}
+            Participants: {self.p1_name} and {self.p2_name}
+            {self.p1_name}'s goal: {self.p1_goal}
+            {self.p2_name}'s goal: {self.p2_goal}
+            """
+            )
+
+
+class ScriptEnvironmentResponse(Message):
+    terminated: bool = Field(
+        description="whether the conversation is terminated",
+        default_factory=lambda: False,
+    )
+    p1_rate: float | tuple[float, dict[str, float]] | None = Field(
+        description="rating of participant 1, on the scale of 1 to 10"
+    )
+    p2_rate: float | tuple[float, dict[str, float]] | None = Field(
+        description="rating of participant 2, on the scale of 1 to 10"
+    )
+    comments: str | None = Field(
+        description="All of the comments supporting the termination and rating"
+    )
+
+    def to_natural_language(self) -> str:
+        reason_to_stop = format_docstring(
+            f"""Environment response:
+        {"The conversation is terminated." if self.terminated else ""}
+        {"Rating of participant 1" + str(self.p1_rate) if self.p1_rate is not None else ""}
+        {"Rating of participant 2" + str(self.p2_rate) if self.p2_rate is not None else ""}
+        {self.comments if self.comments is not None else ""}
+        """
+        )
+        clean_text = ""
+        for line in reason_to_stop.split("\n"):
+            if line.strip():
+                clean_text += line + "\n"
+        return clean_text
+
+
+class AgentAction(Message):
+    action_type: ActionType = Field(
+        description="whether to speak at this turn or choose to not do anything"
+    )
+    argument: str = Field(
+        description="the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action"
+    )
+
+    def to_natural_language(self) -> str:
+        match self.action_type:
+            case "none":
+                return "did nothing"
+            case "speak":
+                return f'said: "{self.argument}"'
+            case "non-verbal communication":
+                return f"[{self.action_type}] {self.argument}"
+            case "action":
+                return f"[{self.action_type}] {self.argument}"
+            case "leave":
+                return "left the conversation"
+
+
+ScriptInteractionReturnType = tuple[
+    list[list[tuple[str, str, Message]]], list[tuple[str, Message]]
+]
+
+
+class ScriptInteraction(Message):
+    interactions: str = Field(
+        description="""The interaction between the two participants in maximum 20 turns. Each turn is separated by a newline, and should only describe one agent. Following the structure:
+        Turn #x
+        [participant's name] [action] {argument for some actions}
+
+        You can use different types of actions, but only use one in each turn. You should move other information into argument part. Below shows a python code snippet of the format for each action type:
+        match self.action_type:
+            case "none":
+                return "did nothing"
+            case "speak":
+                return f'said: "{self.argument}"'
+            case "non-verbal communication":
+                return f"[{self.action_type}] {self.argument}"
+            case "action":
+                return f"[{self.action_type}] {self.argument}"
+            case "leave":
+                return "left the conversation"
+
+        For example, the following is acceptable:
+        Turn #x
+        Oliver Thompson said: "Hey Esmeralda, what's wrong? You seem upset."
+        Turn #x
+        Esmeralda Solis [action] moved closer
+        Turn #x
+        Oliver Thompson [non-verbal communication] smiled
+        Turn #x
+        Esmeralda Solis did nothing
+        Turn #x
+        Oliver Thompson left the conversation
+        Turn #x
+        Esmeralda Solis [action] leaned in and lowered her voice: "Sorry"
+
+        And the following is not acceptable:
+        Turn #1
+        Oliver Thompson [speak] said: "Hey Esmeralda, what's wrong? You seem upset."
+        Turn #1
+        Esmeralda Solis non-verbal communication moved closer
+        """
+    )
+
+    def to_natural_language(self) -> str:
+        return self.interactions
+
+    def parse(
+        self, agent_names: list[str], background: str
+    ) -> tuple[list[list[tuple[str, str, Message]]], list[tuple[str, Message]]]:
+        interaction = self.interactions
+        # print("Interaction: ", interaction)
+        lines = self.split_by_turn(interaction)
+
+        agent_results = []
+        results: list[list[tuple[str, str, Message]]] = [
+            [
+                (
+                    "Environment",
+                    name,
+                    Observation(
+                        last_turn=background,
+                        turn_number=0,
+                        available_actions=["none"],
+                    ),
+                )
+                for name in agent_names
+            ]
+        ]
+
+        for line_idx, line in enumerate(lines):
+            try:
+                res = self.parse_single_dialogue(line)
+                action: AgentAction = cast(AgentAction, res["action"])
+                argument: str = cast(str, res["argument"])
+                cast(int, res["turn"])
+                name: str = cast(str, res["name"])
+
+                parsed_action = AgentAction(action_type=action, argument=argument)
+                if name not in agent_names:
+                    print(
+                        f"The name of the agent, {name}, is not in the list of agent names, {agent_names}"
+                    )
+                    name = agent_names[
+                        line_idx % 2
+                    ]  # TODO Not sure what name to be set here
+            except Exception as e:
+                print(
+                    f"Error when parsing the dialogue: {line}",
+                    f"The error is: {e}",
+                )
+                raise e
+                parsed_action = AgentAction(action_type="none", argument="")
+                name = agent_names[line_idx % 2]  # TODO same question as above
+            inactive_agent_name = (
+                agent_names[0] if name == agent_names[1] else agent_names[1]
+            )
+            results.append(
+                [
+                    (
+                        "Environment",
+                        name,
+                        Observation(
+                            last_turn="environment is the agent",
+                            turn_number=line_idx + 1,
+                            available_actions=["none"],
+                        ),
+                    )
+                    for name in agent_names
+                ]
+                + [
+                    (name, "Environment", parsed_action),
+                    (
+                        inactive_agent_name,
+                        "Environment",
+                        AgentAction(action_type="none", argument="did nothing"),
+                    ),
+                ]
+            )
+
+            agent_results.append((name, parsed_action))
+        # print("Parsed agent results: ", agent_results)
+        return (results, agent_results)  # type: ignore
+
+    def parse_single_dialogue(
+        self, dialogue: str
+    ) -> dict[str, str | int | AgentAction | None]:
+        """Parse a single dialogue string and return a dictionary with turn, name, action, and argument."""
+
+        # Match the turn number and name. Assume all agent name starts with a capital letter and is followed by lowercase letters
+        match_turn_name = re.match(
+            r"Turn #?(\d+):?\s*\n((?:[A-Z]['a-z]* ?)+)", dialogue
+        )
+
+        if not match_turn_name:
+            raise ValueError(
+                f"The dialogue does not match the expected format: {dialogue}"
+            )
+            return None  # TODO Which should we use, return None or raise error?
+
+        turn, name = match_turn_name.groups()
+        action_content = dialogue[
+            len(match_turn_name.group(0)) :
+        ].strip()  # Extract the action content
+
+        # Check for different action types
+        if "did nothing" in action_content:
+            action, argument = "none", ""
+        elif match := re.match(r'said: "(.*?)"', action_content):
+            action, argument = "speak", match.group(1)
+            action, argument = action.strip(), argument.strip()
+        elif match := re.match(r'\[speak\] said: "(.*?)"', action_content):
+            action, argument = "speak", match.group(1)
+            action, argument = action.strip(), argument.strip()
+        elif match := re.match(
+            r"\[(non-verbal communication|action)\] (.*)", action_content
+        ):
+            action, argument = match.groups()
+        elif "left the conversation" in action_content:
+            # TODO Make it more elegant to handle the situation of `left the conversation.`
+            action, argument = "leave", ""
+        else:
+            action, argument = None, None
+
+        parsed_item = {
+            "turn": int(turn),
+            "name": name.strip(),
+            "action": action,
+            "argument": argument,
+        }
+        return parsed_item
+
+    def split_by_turn(self, input_string: str) -> list[str]:
+        """Split the input dialogue string by turn and return a list of dialogues."""
+        # Split using 'Turn #' as delimiter, but keep the delimiter in the results
+        dialogues = re.split(r"(?=Turn #?\d+)", input_string)
+        # Remove any empty strings and strip whitespace
+        dialogues = [dialogue.strip() for dialogue in dialogues if dialogue.strip()]
+        dialogues = [dialogue for dialogue in dialogues if dialogue.startswith("Turn")]
+        # Change from Turn #x to Turn (#)x (# is optional)
+        dialogues[-1] = "\n".join(
+            dialogues[-1].split("\n")[:2]
+        )  # Discard further input in the last turn
+
+        for dialogue in dialogues:
+            # TODO this is current workaround for the issue of multiple agents in one turn
+            if len(dialogue.split("\n")) >= 3:
+                raise ValueError("Only one agent can act per turn.")
+        return dialogues
+
+    @staticmethod
+    def default_value_for_return_type() -> ScriptInteractionReturnType:
+        results_1: list[list[tuple[str, str, Message]]] = [
+            [
+                (
+                    "Environment",
+                    name,
+                    Observation(
+                        last_turn="Environment is the agent",
+                        turn_number=0,
+                        available_actions=["none"],
+                    ),
+                )
+                for name in ["none", "none"]
+            ]
+        ]
+        results_2: list[tuple[str, Message]] = [
+            ("", AgentAction(action_type="none", argument=""))
+        ]
+        return (results_1, results_2)

requirements.txt

@@ -1,6 +1,148 @@
-sotopia
-gradio
-transformers
-torch
-peft
-bitsandbytes
+absl-py==2.1.0
+accelerate==0.29.3
+aiofiles==23.2.1
+aiohttp==3.9.5
+aiosignal==1.3.1
+altair==5.3.0
+annotated-types==0.6.0
+anyio==3.7.1
+attrs==23.2.0
+beartype==0.14.1
+bitsandbytes==0.43.1
+certifi==2024.2.2
+cffi==1.16.0
+charset-normalizer==3.3.2
+click==8.1.7
+cloudpickle==3.0.0
+contourpy==1.2.1
+cryptography==42.0.5
+cycler==0.12.1
+dataclasses-json==0.6.4
+datasets==2.18.0
+dill==0.3.8
+distro==1.9.0
+Farama-Notifications==0.0.4
+ffmpy==0.3.2
+filelock==3.13.4
+fonttools==4.51.0
+frozenlist==1.4.1
+fsspec==2024.2.0
+gin-config==0.5.0
+gradio==4.27.0
+gradio_client==0.15.1
+greenlet==3.0.3
+gymnasium==0.29.1
+h11==0.14.0
+hiredis==2.3.2
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.22.2
+idna==3.7
+importlib_metadata==7.1.0
+importlib_resources==6.4.0
+Jinja2==3.1.3
+jsonpatch==1.33
+jsonpointer==2.4
+jsonschema==4.21.1
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+langchain==0.1.16
+langchain-community==0.0.33
+langchain-core==0.1.44
+langchain-openai==0.0.5
+langchain-text-splitters==0.0.1
+langsmith==0.1.48
+litellm==1.35.12
+lxml==4.9.4
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+marshmallow==3.21.1
+matplotlib==3.8.4
+mdurl==0.1.2
+more-itertools==10.2.0
+mpmath==1.3.0
+multidict==6.0.5
+multiprocess==0.70.16
+mypy==1.9.0
+mypy-extensions==1.0.0
+names==0.3.0
+networkx==3.3
+numpy==1.26.4
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-nccl-cu12==2.19.3
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.1.105
+openai==1.22.0
+orjson==3.10.1
+packaging==23.2
+pandas==2.2.2
+pandas-stubs==2.2.1.240316
+peft==0.10.0
+pettingzoo==1.24.0
+pillow==10.3.0
+psutil==5.9.8
+pyarrow==15.0.2
+pyarrow-hotfix==0.6
+pycparser==2.22
+pydantic==2.7.0
+pydantic_core==2.18.1
+pydub==0.25.1
+Pygments==2.17.2
+pyparsing==3.1.2
+python-dateutil==2.9.0.post0
+python-dotenv==1.0.1
+python-multipart==0.0.9
+python-ulid==1.1.0
+pytz==2024.1
+PyYAML==6.0.1
+redis==5.0.3
+referencing==0.34.0
+regex==2024.4.16
+requests==2.31.0
+rich==13.7.1
+rpds-py==0.18.0
+ruff==0.3.7
+safetensors==0.4.3
+scipy==1.13.0
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+SQLAlchemy==2.0.29
+sseclient-py==1.8.0
+starlette==0.27.0
+sympy==1.12
+tabulate==0.9.0
+tenacity==8.2.3
+tiktoken==0.5.2
+tokenizers==0.19.1
+tomlkit==0.12.0
+toolz==0.12.1
+torch==2.2.2
+tqdm==4.66.2
+transformers==4.40.0
+triton==2.2.0
+typer==0.12.3
+types-cffi==1.16.0.20240331
+types-pyOpenSSL==24.0.0.20240417
+types-pytz==2024.1.0.20240417
+types-redis==4.6.0.20240417
+types-setuptools==69.5.0.20240415
+types-tqdm==4.66.0.20240417
+typing-inspect==0.9.0
+typing_extensions==4.11.0
+tzdata==2024.1
+urllib3==2.2.1
+uvicorn==0.23.2
+websockets==11.0.3
+xxhash==3.4.1
+yarl==1.9.4
+zipp==3.18.1

sotopia_pi_generate.py

@@ -1,17 +1,18 @@
 import re
+from typing import TypeVar
+from functools import cache
+import logging
 
 import torch
-from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
 )
-
+from peft import PeftModel
 from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from langchain_community.chat_models import ChatLiteLLM
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-
 from langchain.chains import LLMChain
 from langchain.output_parsers import PydanticOutputParser
 from langchain.prompts import (
@@ -20,17 +21,16 @@ from langchain.prompts import (
     PromptTemplate,
 )
 from langchain.schema import BaseOutputParser, OutputParserException
-from typing import TypeVar
+from message_classes import ActionType, AgentAction
+from utils import format_docstring
 
-from sotopia.messages import ActionType, AgentAction
-from sotopia.utils import format_docstring
-from functools import cache
-import logging
+from langchain_callback_handler import LoggingCallbackHandler
 
-OutputType = TypeVar("OutputType", bound=object)
+HF_TOKEN_KEY_FILE="./hf_token.key"
 
+OutputType = TypeVar("OutputType", bound=object)
 log = logging.getLogger("generate")
-# logging_handler = LoggingCallbackHandler("langchain")
+logging_handler = LoggingCallbackHandler("langchain")
 
 def generate_action(
     model_name: str,
@@ -39,7 +39,7 @@ def generate_action(
     action_types: list[ActionType],
     agent: str,
     temperature: float = 0.7,
-) -> tuple[AgentAction, str]:
+) -> AgentAction:
     """
     Using langchain to generate an example episode
     """
@@ -73,14 +73,26 @@ def generate_action(
             temperature=temperature,
         )
     except Exception:
-        return AgentAction(action_type="none", argument=""), ""
+        return AgentAction(action_type="none", argument="")
 
 @cache
-def prepare_model(model_name):
+def prepare_model(model_name, hf_token_key_file=HF_TOKEN_KEY_FILE):
     compute_type = torch.float16
+    with open (hf_token_key_file, 'r') as f:
+        hf_token = f.read().strip()
     
-    if 'cmu-lti/sotopia-pi-mistral-7b-BC_SR'in model_name:
-        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", token="REDACTED")
+    if model_name == 'cmu-lti/sotopia-pi-mistral-7b-BC_SR':
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", token=hf_token)
+        model = AutoModelForCausalLM.from_pretrained(
+        "mistralai/Mistral-7B-Instruct-v0.1",
+        cache_dir="./.cache",
+        device_map='cuda',
+        token=hf_token
+        )
+        model = PeftModel.from_pretrained(model, model_name).to("cuda")
+        
+    elif model_name == 'cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit':
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", token=hf_token)
         model = AutoModelForCausalLM.from_pretrained(
         "mistralai/Mistral-7B-Instruct-v0.1",
         cache_dir="./.cache",
@@ -91,11 +103,22 @@ def prepare_model(model_name):
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=compute_type,
             ),
-        token="REDACTED"
+        token=hf_token
         )
         model = PeftModel.from_pretrained(model, model_name).to("cuda")
+    
+    elif model_name == 'mistralai/Mistral-7B-Instruct-v0.1':
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", token=hf_token)
+        model = AutoModelForCausalLM.from_pretrained(
+        "mistralai/Mistral-7B-Instruct-v0.1",
+        cache_dir="./.cache",
+        device_map='cuda',
+        token=hf_token
+        )
+        
     else:
          raise RuntimeError(f"Model {model_name} not supported")
+     
     return model, tokenizer
 
 def obtain_chain_hf(
@@ -111,9 +134,17 @@ def obtain_chain_hf(
     )
     chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])
     model, tokenizer = prepare_model(model_name)
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=max_tokens, temperature=temperature)
+    pipe = pipeline("text-generation", 
+                    model=model, 
+                    tokenizer=tokenizer, 
+                    max_new_tokens=100, 
+                    temperature=temperature, 
+                    return_full_text=False, 
+                    do_sample=True,
+                    num_beams=3,
+                    length_penalty=-1.0,
+                    )
     hf = HuggingFacePipeline(pipeline=pipe)
-    # import pdb; pdb.set_trace()
     chain = LLMChain(llm=hf, prompt=chat_prompt_template)
     return chain
 
@@ -123,7 +154,7 @@ def generate(
     input_values: dict[str, str],
     output_parser: BaseOutputParser[OutputType],
     temperature: float = 0.7,
-) -> tuple[OutputType, str]:
+) -> OutputType:
     # import pdb; pdb.set_trace()
     input_variables = re.findall(r"{(.*?)}", template)
     assert (
@@ -135,8 +166,9 @@ def generate(
     chain = obtain_chain(model_name, template, input_variables, temperature)
     if "format_instructions" not in input_values:
         input_values["format_instructions"] = output_parser.get_format_instructions()
-    result = chain.predict([], **input_values)
-    # import pdb; pdb.set_trace()
+    result = chain.predict([logging_handler], **input_values)
+    prompt = logging_handler.retrive_prompt()
+    import pdb; pdb.set_trace()
     try:
         parsed_result = output_parser.parse(result)
     except KeyboardInterrupt:
@@ -146,6 +178,7 @@ def generate(
             f"[red] Failed to parse result: {result}\nEncounter Exception {e}\nstart to reparse",
             extra={"markup": True},
         )
+        import pdb; pdb.set_trace()
         reformat_parsed_result = format_bad_output(
             result, format_instructions=output_parser.get_format_instructions()
         )
@@ -175,7 +208,7 @@ def format_bad_output(
         "ill_formed_output": ill_formed_output,
         "format_instructions": format_instructions,
     }
-    reformat = chain.predict([], **input_values)
+    reformat = chain.predict([logging_handler], **input_values)
     log.info(f"Reformated output: {reformat}")
     return reformat
 
@@ -189,7 +222,7 @@ def obtain_chain(
     """
     Using langchain to sample profiles for participants
     """
-    if model_name in ["cmu-lti/sotopia-pi-mistral-7b-BC_SR"]:
+    if model_name in ["cmu-lti/sotopia-pi-mistral-7b-BC_SR", "cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit"]:
         return obtain_chain_hf(
             model_name=model_name,
             template=template,
@@ -212,32 +245,6 @@ def obtain_chain(
     chain = LLMChain(llm=chat, prompt=chat_prompt_template)
     return chain
 
-def format_bad_output(
-    ill_formed_output: str,
-    format_instructions: str,
-    model_name: str = "gpt-3.5-turbo",
-) -> str:
-    template = """
-    Given the string that can not be parsed by json parser, reformat it to a string that can be parsed by json parser.
-    Original string: {ill_formed_output}
-
-    Format instructions: {format_instructions}
-
-    Please only generate the JSON:
-    """
-    chain = obtain_chain(
-        model_name=model_name,
-        template=template,
-        input_variables=re.findall(r"{(.*?)}", template),
-    )
-    input_values = {
-        "ill_formed_output": ill_formed_output,
-        "format_instructions": format_instructions,
-    }
-    reformat = chain.predict([], **input_values)
-    log.info(f"Reformated output: {reformat}")
-    return reformat
-
 def _return_fixed_model_version(model_name: str) -> str:
     return {
         "gpt-3.5-turbo": "gpt-3.5-turbo-0613",

utils.py

@@ -1,5 +1,6 @@
 from typing import List, Tuple
 import ast
+import re
 
 class Agent:
     def __init__(self, agent_profile):
@@ -31,80 +32,10 @@ class Environment:
         self.agent_goals = env_profile["agent_goals"]
         self.relationship = env_profile["relationship"]
         
-def get_format_guide():
-    return """ Your available action types are
-    "none action speak non-verbal communication leave".
-    Note: You can "leave" this conversation if 1. you have achieved your social goals, 2. this conversation makes you uncomfortable, 3. you find it uninteresting/you lose your patience, 4. or for other reasons you want to leave.
-
-    Please only generate a JSON string including the action type and the argument.
-    Your action should follow the given format:
-    \nAs an example, for the schema {\"properties\": {\"foo\": {\"title\": \"Foo\", \"description\": \"a list of strings\", \"type\": \"array\", \"items\": {\"type\": \"string\"}}}, \"required\": [\"foo\"]}
-    the object {\"foo\": [\"bar\", \"baz\"]} is a well-formatted instance of the schema. The object {\"properties\": {\"foo\": [\"bar\", \"baz\"]}} is not well-formatted.
-    \nHere is the output schema:\n```\n{\"description\": \"An interface for messages.\\nThere is only one required method: to_natural_language\", \"properties\": {\"action_type\": {\"title\": \"Action Type\", \"description\": \"whether to speak at this turn or choose to not do anything\", \"enum\": [\"none\", \"speak\", \"non-verbal communication\", \"action\", \"leave\"], \"type\": \"string\"}, \"argument\": {\"title\": \"Argument\", \"description\": \"the utterance if choose to speak, the expression or gesture if choose non-verbal communication, or the physical action if choose action\", \"type\": \"string\"}}, \"required\": [\"action_type\", \"argument\"]}\n```\u001b[0m
-    """
-
-def get_starter_prompt(machine_agent, human_agent, environment):
-    return f"Imagine you are {machine_agent.name}, your task is to act/speak as {machine_agent.name} would, keeping in mind {machine_agent.name}'s social goal.\nYou can find {machine_agent.name}'s background and goal in the 'Here is the context of the interaction' field.\nNote that {machine_agent.name}'s secret and goal is only visible to you.\nYou should try your best to achieve {machine_agent.name}'s goal in a way that align with their character traits.\nAdditionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before).\n\nHere is the context of this interaction:\n Scenario: {environment.scenario}\nParticipants: {human_agent.name} and {machine_agent.name}\n{human_agent.name}'s background: {human_agent.background} Personality and values description: {human_agent.personality} \n{machine_agent.name}'s background: {machine_agent.background} Personality and values description: {machine_agent.personality} {machine_agent.name}'s secrets: {machine_agent.secret}\n{human_agent.name}'s goal: Unknown\n{machine_agent.name}'s goal: {environment.agent_goals[1]}\nConversation Starts:"
-
+        
 def get_context_prompt(machine_agent, human_agent, environment):
     return f"Here is the context of this interaction:\n Scenario: {environment.scenario}\nParticipants: {human_agent.name} and {machine_agent.name}\n{human_agent.name}'s background: {human_agent.background} Personality and values description: {human_agent.personality} \n{machine_agent.name}'s background: {machine_agent.background} Personality and values description: {machine_agent.personality} {machine_agent.name}'s secrets: {machine_agent.secret}\n{human_agent.name}'s goal: Unknown\n{machine_agent.name}'s goal: {environment.agent_goals[1]}\nConversation Starts:"
-
-
-# we define history as
-# [(user_message, bot_message), (user_message, bot_message)]
-
-# we define dialogue history as
-# user_name: user_message\nbot_name: bot_message\nuser_name: user_message\nbot_name: bot_message\n
-
-
-def dialogue_history_length_check(string, max_token, tokenizer):
-    prompt_tokens = len(tokenizer(string)["input_ids"])
-    return max(prompt_tokens - max_token, 0)
-
-
-def truncate_dialogue_history_to_length(dia_his, surpass_num, tokenizer):
-    dia_sen = dia_his.split("\n")
-    remove_len = 0
-    i = 0
-    while remove_len < surpass_num:
-        remove_len += len(tokenizer(dia_sen[i])["input_ids"])
-        i += 1
-    trunc_dia = "\n".join(p for p in dia_sen[i:])
-    return trunc_dia
-
-
-def format_bot_message(bot_message) -> str:
-    # # import pdb; pdb.set_trace()
-    start_idx, end_idx = bot_message.index("{"), bot_message.index("}")
-    if end_idx == -1:
-        bot_message += "'}"
-        end_idx = len(bot_message)
-    json_response = ast.literal_eval(bot_message[start_idx:end_idx+1])
-    match json_response["action_type"]:
-        case "none":
-            return 'did nothing'
-        case "speak":
-            return json_response["argument"]
-        case "non-verbal communication":
-            return f'[{json_response["action_type"]}] {json_response["argument"]}'
-        case "action":
-            return f'[{json_response["action_type"]}] {json_response["argument"]}'
-        case "leave":
-            return 'left the conversation'
-    
-def dialogue_history_creation(history, user_name, bot_name):
-    dialogue_history = ""
-    for idx, turn in enumerate(history):
-        user_message, bot_message = turn
-        # TODOTODO (haofeiyu): we first assume that human talks first
-        user_turn_idx = idx * 2
-        bot_turn_idx = idx * 2 + 1
-        if not bot_message.startswith("["): # if action type == speak, need to add 'said: ' to be consistent with the dialog prompt
-            bot_message = "said :" + bot_message
-        dialogue_history = f"{dialogue_history}\n\nTurn #{user_turn_idx}: {user_name}: {user_message}\n\nTurn #{bot_turn_idx}: {bot_name}: {bot_message}"
-    last_turn_idx = len(history) * 2
-    return dialogue_history, last_turn_idx
-
+        
 def dialogue_history_prompt(message, history, user_agent, bot_agent):
     dialogue_history = ""
     for idx, turn in enumerate(history):
@@ -117,31 +48,8 @@ def dialogue_history_prompt(message, history, user_agent, bot_agent):
         dialogue_history = f"{dialogue_history}\n\nTurn #{user_turn_idx}: {user_agent.name}: {user_message}\n\nTurn #{bot_turn_idx}: {bot_agent.name}: {bot_message}"
     last_turn_idx = len(history) * 2
     dialogue_history = f"{dialogue_history}\n\nTurn #{last_turn_idx+1}: {user_agent.name}: {message}\n."
-    return dialogue_history, last_turn_idx+2
-
-
-def dialogue_history_truncation(dialogue_history, max_token_num, tokenizer):
-    surpass_num = dialogue_history_length_check(
-        dialogue_history, max_token_num, tokenizer
-    )
-    if surpass_num > 0:
-        dialogue_history = truncate_dialogue_history_to_length(
-            dialogue_history, surpass_num, tokenizer
-        )
-    return dialogue_history
-
+    return dialogue_history, last_turn_idx + 2
 
-def format_hostory_prompt(
-    message: str,
-    history: List[Tuple[str, str]],
-    instructions: str,
-    user_name: str,
-    bot_name: str,
-) -> str:
-    prompt = instructions.strip()
-    dialogue_history, last_turn_idx = dialogue_history_creation(
-        history, user_name, bot_name
-    )
-    prompt = f"{prompt}\n{dialogue_history}"
-    prompt = f"{prompt}\n\nTurn #{last_turn_idx+1}: {user_name}: {message}\n.\nYou are at Turn #{last_turn_idx+2}."
-    return prompt
+def format_docstring(docstring: str) -> str:
+    """Format a docstring for use in a prompt template."""
+    return re.sub("\n +", "\n", docstring).strip()