fixed parsing errors and extra_info (#56)

app.py

@@ -1,6 +1,7 @@
 import os
 from collections import defaultdict
 import json
+from typing import Literal
 
 import gradio as gr
 
@@ -25,6 +26,21 @@ RELATIONSHIP_PROFILES = "profiles/relationship_profiles.jsonl"
 
 ACTION_TYPES = ['none', 'action', 'non-verbal communication', 'speak', 'leave']
 
+MODEL_OPTIONS = [
+    "gpt-3.5-turbo",
+    "gpt-4",
+    "gpt-4-turbo",
+    "cmu-lti/sotopia-pi-mistral-7b-BC_SR",
+    "cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit",
+    "mistralai/Mistral-7B-Instruct-v0.1"
+    # "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    # "togethercomputer/llama-2-7b-chat",
+    # "togethercomputer/llama-2-70b-chat",
+    # "togethercomputer/mpt-30b-chat",
+    # "together_ai/togethercomputer/llama-2-7b-chat",
+    # "together_ai/togethercomputer/falcon-7b-instruct",
+]
+
 @cache
 def get_sotopia_profiles(env_file=ENVIRONMENT_PROFILES, agent_file=AGENT_PROFILES, relationship_file=RELATIONSHIP_PROFILES):
     with open(env_file, 'r') as f:
@@ -126,13 +142,27 @@ def create_bot_info(bot_agent_dropdown):
     return gr.Textbox(label="Bot Agent Profile", lines=4, value=text)
 
 def create_user_goal(environment_dropdown):
-     _, environment_dict, _, _ = get_sotopia_profiles()
-     text = environment_dict[environment_dropdown].agent_goals[0]
-     return gr.Textbox(label="User Agent Goal", lines=4, value=text)
+    _, environment_dict, _, _ = get_sotopia_profiles()
+    text = environment_dict[environment_dropdown].agent_goals[0]
+    text = text.replace('(', '').replace(')', '')
+    if "<extra_info>" in text:
+        text = text.replace("<extra_info>", "\n\n")
+        text = text.replace("</extra_info>", "\n")
+    if "<strategy_hint>" in text:
+        text = text.replace("<strategy_hint>", "\n\n")
+        text = text.replace("</strategy_hint>", "\n")
+    return gr.Textbox(label="User Agent Goal", lines=4, value=text)
 
 def create_bot_goal(environment_dropdown):
     _, environment_dict, _, _ = get_sotopia_profiles()
     text = environment_dict[environment_dropdown].agent_goals[1]
+    text = text.replace('(', '').replace(')', '')
+    if "<extra_info>" in text:
+        text = text.replace("<extra_info>", "\n\n")
+        text = text.replace("</extra_info>", "\n")
+    if "<strategy_hint>" in text:
+        text = text.replace("<strategy_hint>", "\n\n")
+        text = text.replace("</strategy_hint>", "\n")
     return gr.Textbox(label="Bot Agent Goal", lines=4, value=text)
 
 def sotopia_info_accordion(accordion_visible=True):
@@ -147,7 +177,7 @@ def sotopia_info_accordion(accordion_visible=True):
                 interactive=True,
             )
             model_name_dropdown = gr.Dropdown(
-                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"],
+                choices=MODEL_OPTIONS,
                 value=DEFAULT_MODEL_SELECTION,
                 interactive=True,
                 label="Model Selection"
@@ -215,7 +245,7 @@ def chat_tab():
         
         context = get_context_prompt(bot_agent, user_agent, environment)
         dialogue_history, next_turn_idx = dialogue_history_prompt(message, history, user_agent, bot_agent)
-        prompt_history = f"{context}\n\n{dialogue_history}"
+        prompt_history = f"{context}{dialogue_history}"
         agent_action = generate_action(model_selection, prompt_history, next_turn_idx, ACTION_TYPES, bot_agent.name, TEMPERATURE)
         return agent_action.to_natural_language()
     

message_classes.py

@@ -120,7 +120,7 @@ class AgentAction(Message):
             case "none":
                 return "did nothing"
             case "speak":
-                return f'said: "{self.argument}"'
+                return f"{self.argument}"
             case "non-verbal communication":
                 return f"[{self.action_type}] {self.argument}"
             case "action":

sotopia_pi_generate.py

@@ -28,6 +28,9 @@ from utils import format_docstring
 from langchain_callback_handler import LoggingCallbackHandler
 
 HF_TOKEN_KEY_FILE="./hf_token.key"
+if os.path.exists(HF_TOKEN_KEY_FILE):
+    with open(HF_TOKEN_KEY_FILE, "r") as f:
+        os.environ["HF_TOKEN"] = f.read().strip()
 
 OutputType = TypeVar("OutputType", bound=object)
 log = logging.getLogger("generate")
@@ -44,59 +47,54 @@ def generate_action(
     """
     Using langchain to generate an example episode
     """
-    try:
+    # try:
         # Normal case, model as agent
-        template = """
-            Imagine you are {agent}, your task is to act/speak as {agent} would, keeping in mind {agent}'s social goal.
-            You can find {agent}'s goal (or background) in the 'Here is the context of the interaction' field.
-            Note that {agent}'s goal is only visible to you.
-            You should try your best to achieve {agent}'s goal in a way that align with their character traits.
-            Additionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before).
-            {history}.
-            You are at Turn #{turn_number}. Your available action types are
-            {action_list}.
-            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.
+    template = """
+        Imagine you are {agent}, your task is to act/speak as {agent} would, keeping in mind {agent}'s social goal.
+        You can find {agent}'s goal (or background) in the 'Here is the context of the interaction' field.
+        Note that {agent}'s goal is only visible to you.
+        You should try your best to achieve {agent}'s goal in a way that align with their character traits.
+        Additionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before).\n
+        {history}.
+        You are at Turn #{turn_number}. Your available action types are
+        {action_list}.
+        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:
-            {format_instructions}
-        """
-        return generate(
-            model_name=model_name,
-            template=template,
-            input_values=dict(
-                agent=agent,
-                turn_number=str(turn_number),
-                history=history,
-                action_list=" ".join(action_types),
-            ),
-            output_parser=PydanticOutputParser(pydantic_object=AgentAction),
-            temperature=temperature,
-        )
-    except Exception:
-        return AgentAction(action_type="none", argument="")
+        Please only generate a JSON string including the action type and the argument.
+        Your action should follow the given format:
+        {format_instructions}
+    """
+    return generate(
+        model_name=model_name,
+        template=template,
+        input_values=dict(
+            agent=agent,
+            turn_number=str(turn_number),
+            history=history,
+            action_list=" ".join(action_types),
+        ),
+        output_parser=PydanticOutputParser(pydantic_object=AgentAction),
+        temperature=temperature,
+    )
+    # except Exception as e:
+    #     print(e)
+    #     return AgentAction(action_type="none", argument="")
 
 @cache
-def prepare_model(model_name, hf_token_key_file=HF_TOKEN_KEY_FILE):
+def prepare_model(model_name):
     compute_type = torch.float16
-    if os.path.exists(hf_token_key_file):
-        with open (hf_token_key_file, 'r') as f:
-            hf_token = f.read().strip()
-    else:
-        hf_token = os.environ["HF_TOKEN"]
     
     if model_name == 'cmu-lti/sotopia-pi-mistral-7b-BC_SR':
-        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", token=hf_token)
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", model_max_length=4096)
         model = AutoModelForCausalLM.from_pretrained(
         "mistralai/Mistral-7B-Instruct-v0.1",
         cache_dir="./.cache",
-        device_map='cuda',
-        token=hf_token
+        device_map='cuda'
         )
         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)
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", model_max_length=4096)
         model = AutoModelForCausalLM.from_pretrained(
         "mistralai/Mistral-7B-Instruct-v0.1",
         cache_dir="./.cache",
@@ -106,18 +104,17 @@ def prepare_model(model_name, hf_token_key_file=HF_TOKEN_KEY_FILE):
             bnb_4bit_use_double_quant=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype=compute_type,
-            ),
-        token=hf_token
+            )
         )
-        model = PeftModel.from_pretrained(model, model_name).to("cuda")
+        model = PeftModel.from_pretrained(model, model_name[0:-5]).to("cuda")
     
     elif model_name == 'mistralai/Mistral-7B-Instruct-v0.1':
-        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", token=hf_token)
+        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", model_max_length=4096)
+        tokenizer.model_max_length = 4096
         model = AutoModelForCausalLM.from_pretrained(
         "mistralai/Mistral-7B-Instruct-v0.1",
         cache_dir="./.cache",
-        device_map='cuda',
-        token=hf_token
+        device_map='cuda'
         )
         
     else:
@@ -146,7 +143,6 @@ def obtain_chain_hf(
                     return_full_text=False, 
                     do_sample=True,
                     num_beams=3,
-                    length_penalty=-1.0,
                     )
     hf = HuggingFacePipeline(pipeline=pipe)
     chain = LLMChain(llm=hf, prompt=chat_prompt_template)
@@ -171,6 +167,8 @@ def generate(
         input_values["format_instructions"] = output_parser.get_format_instructions()
     result = chain.predict([logging_handler], **input_values)
     prompt = logging_handler.retrive_prompt()
+    print(f"Prompt:\n {prompt}")
+    print(f"Result:\n {result}")
     try:
         parsed_result = output_parser.parse(result)
     except KeyboardInterrupt:
@@ -183,6 +181,7 @@ def generate(
         reformat_parsed_result = format_bad_output(
             result, format_instructions=output_parser.get_format_instructions()
         )
+        print(f"Reformatted result:\n {reformat_parsed_result}")
         parsed_result = output_parser.parse(reformat_parsed_result)
     log.info(f"Generated result: {parsed_result}")
     return parsed_result
@@ -223,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", "cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit"]:
+    if model_name in ["cmu-lti/sotopia-pi-mistral-7b-BC_SR", "cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit", "mistralai/Mistral-7B-Instruct-v0.1"]:
         return obtain_chain_hf(
             model_name=model_name,
             template=template,
@@ -247,10 +246,11 @@ def obtain_chain(
     return chain
 
 def _return_fixed_model_version(model_name: str) -> str:
-    return {
+    model_version_map = {
         "gpt-3.5-turbo": "gpt-3.5-turbo-0613",
         "gpt-3.5-turbo-finetuned": "ft:gpt-3.5-turbo-0613:academicscmu::8nY2zgdt",
         "gpt-3.5-turbo-ft-MF": "ft:gpt-3.5-turbo-0613:academicscmu::8nuER4bO",
         "gpt-4": "gpt-4-0613",
         "gpt-4-turbo": "gpt-4-1106-preview",
-    }[model_name]
+    }
+    return model_version_map[model_name] if model_name in model_version_map else model_name

start_app.sh

@@ -0,0 +1,4 @@
+export OPENAI_API_KEY=$(cat openai_api.key)
+export HF_TOKEN=$(cat hf_token.key)
+
+python app.py

utils.py

@@ -44,11 +44,11 @@ def dialogue_history_prompt(message, history, user_agent, bot_agent):
         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_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
+            bot_message = 'said:"' + bot_message + '"'
+        dialogue_history = f"""{dialogue_history}\n\nTurn #{user_turn_idx} {user_agent.name} said: "{user_message}"\n\nTurn #{bot_turn_idx}: {bot_agent.name}: {bot_message}"""
+    curr_turn_idx = len(history) * 2
+    dialogue_history = f"""{dialogue_history}\n\nTurn #{curr_turn_idx} {user_agent.name} said: "{message}"\n"""
+    return dialogue_history, curr_turn_idx + 1
 
 def format_docstring(docstring: str) -> str:
     """Format a docstring for use in a prompt template."""