v0.5 (#4)

- v0.5 utilities (87cfe08d079faab66e43edd83437c18cd8f6d2f7)

app_config.py

@@ -4,14 +4,21 @@ from models.model_seeds import seeds, seed2str
 ISSUES = [k for k,_ in seeds.items()]
 SOURCES = [
     "CTL_llama2",
+    "CTL_mistral",
     'OA_rolemodel', 
         #    'OA_finetuned',
 ]
 SOURCES_LAB = {"OA_rolemodel":'OpenAI GPT3.5',
                "OA_finetuned":'Finetuned OpenAI',
-               "CTL_llama2": "Custom CTL"
+               "CTL_llama2": "Llama",
+               "CTL_mistral": "Mistral",
                }
 
+ENDPOINT_NAMES = {
+    'CTL_llama2': "llama2_convo_sim",
+    "CTL_mistral": "convo_sim_mistral"
+}
+
 def source2label(source):
     return SOURCES_LAB[source]
 

convosim.py

@@ -5,11 +5,13 @@ from langchain.schema.messages import HumanMessage
 from utils.mongo_utils import get_db_client
 from utils.app_utils import create_memory_add_initial_message, get_random_name,  DEFAULT_NAMES_DF
 from utils.memory_utils import clear_memory, push_convo2db
-from utils.chain_utils import get_chain
+from utils.chain_utils import get_chain, custom_chain_predict
 from app_config import ISSUES, SOURCES, source2label, issue2label
 
 logger = get_logger(__name__)
 openai_api_key = os.environ['OPENAI_API_KEY'] 
+temperature = 0.8
+username = "barb-chase" #"ivnban-ctl"
 
 if "sent_messages" not in st.session_state:
     st.session_state['sent_messages'] = 0
@@ -28,8 +30,8 @@ if 'texter_name' not in st.session_state:
 memories = {'memory':{"issue": st.session_state['issue'], "source": st.session_state['previous_source']}}
 
 with st.sidebar:
-    username = st.text_input("Username", value='ivnban-ctl', max_chars=30)
-    temperature = st.slider("Temperature", 0., 1., value=0.8, step=0.1)
+    # username = st.text_input("Username", value='ivnban-ctl', max_chars=30)
+    # temperature = st.slider("Temperature", 0., 1., value=0.8, step=0.1)
     issue = st.selectbox("Select a Scenario", ISSUES, index=0, format_func=issue2label,
                             on_change=clear_memory, kwargs={"memories":memories, "username":username, "language":"English"}
                         )
@@ -76,6 +78,8 @@ if prompt := st.chat_input():
         push_convo2db(memories, username, language)
 
     st.chat_message("user").write(prompt)
-    response = llm_chain.predict(input=prompt, stop=stopper)
+    responses = custom_chain_predict(llm_chain, prompt, stopper)
+    # responses = llm_chain.predict(input=prompt, stop=stopper)
     # response = update_memory_completion(prompt, st.session_state["memory"], OA_engine, temperature)
-    st.chat_message("assistant").write(response)
+    for response in responses:
+        st.chat_message("assistant").write(response)

models/custom_parsers.py

@@ -1,5 +1,7 @@
 from typing import List
+import re
 from langchain.schema import BaseOutputParser
+from itertools import chain
 
 class CustomStringOutputParser(BaseOutputParser[List[str]]):
     """Parse the output of an LLM call to a list."""
@@ -10,8 +12,46 @@ class CustomStringOutputParser(BaseOutputParser[List[str]]):
 
     def parse(self, text: str) -> str:
         """Parse the output of an LLM call."""
-        text = text.split("texter:")[0]
         text = text.split("helper")[0]
+        text = text.split("\nhelper")[0]
         text = text.rstrip("\n")
-        text = text.strip()
-        return text
+        text_list = text.split("texter:")
+        text_list = [x.split("\ntexter") for x in text_list]
+        text_list = [x.strip() for x in list(chain.from_iterable(text_list))]
+        return text_list
+    
+class CustomINSTOutputParser(BaseOutputParser[List[str]]):
+    """Parse the output of an LLM call to a list."""
+
+    name = "Kit"
+    name_rx = re.compile(r""+ name + r":|" + name.lower() + r":")
+    whispers = re.compile((r"([\(]).*?([\)])"))
+    reactions = re.compile(r"([\*]).*?([\*])")
+    double_spaces = re.compile(r"  ")
+    quotation_rx = re.compile('"')
+
+    @property
+    def _type(self) -> str:
+        return "str"
+
+    def parse_whispers(self, text: str) -> str:
+        text = self.name_rx.sub("", text).strip()
+        text = self.reactions.sub("", text).strip()
+        text = self.whispers.sub("", text).strip()
+        text = self.double_spaces.sub(r" ", text).strip()
+        text = self.quotation_rx.sub("", text).strip()
+        return text 
+    
+    def parse_split(self, text: str) -> str:
+        text = text.split("[INST]")[0]
+        text_list = text.split("[/INST]")
+        text_list = [x.split("</s>") for x in text_list]
+        text_list = [x.strip() for x in list(chain.from_iterable(text_list))]
+        text_list = [x.split("\n\n") for x in text_list]
+        text_list = [x.strip().rstrip("\n") for x in list(chain.from_iterable(text_list))]
+        return text_list
+    
+    def parse(self, text: str) -> str:
+        text = self.parse_whispers(text)
+        text_list = self.parse_split(text)
+        return text_list

models/databricks/instruction_memory.py

@@ -0,0 +1,129 @@
+import os
+import sys
+from typing import Any, Dict, List, Sequence, Union
+
+from langchain.memory.chat_memory import BaseChatMemory
+from langchain.schema.messages import (
+    AIMessage,
+    BaseMessage,
+    ChatMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+)
+
+# Hacky way of importing other directory modules.
+# This is done this way instead of src.evaluation.postprocessing_utils
+# Because this file is used for inference where there is not the same folder structure
+sys.path.append(os.path.join(os.path.dirname(__file__), "..", "evaluation"))
+
+def get_buffer_string_inst(
+    messages: Sequence[BaseMessage],
+    human_prefix: str = "[INST]",
+    ai_prefix: str = "[/INST]",
+) -> str:
+    """Convert sequence of Messages to strings and concatenate them into one string.
+
+    Args:
+        messages: Messages to be converted to strings.
+        human_prefix: The prefix to prepend to contents of HumanMessages.
+        ai_prefix: THe prefix to prepend to contents of AIMessages.
+
+    Returns:
+        A single string concatenation of all input messages.
+
+    Example:
+        .. code-block:: python
+
+            from langchain.schema import AIMessage, HumanMessage
+
+            messages = [
+                HumanMessage(content="Hi, how are you?"),
+                AIMessage(content="Good, how are you?"),
+            ]
+            get_buffer_string(messages)
+            # -> "Human: Hi, how are you?\nAI: Good, how are you?"
+    """
+    string_messages = []
+    for m in messages:
+        if isinstance(m, HumanMessage):
+            role = human_prefix
+        elif isinstance(m, AIMessage):
+            role = ai_prefix
+        elif isinstance(m, SystemMessage):
+            role = "System"
+        elif isinstance(m, FunctionMessage):
+            role = "Function"
+        elif isinstance(m, ChatMessage):
+            role = m.role
+        else:
+            raise ValueError(f"Got unsupported message type: {m}")
+        message = f"{role} {m.content}"
+        if isinstance(m, AIMessage) and "function_call" in m.additional_kwargs:
+            message += f"{m.additional_kwargs['function_call']}"
+        string_messages.append(message)
+
+    return " ".join(string_messages)
+
+
+class CustomBufferInstructionMemory(BaseChatMemory):
+    memory_key: str = "chat_history"
+    input_key: str = "input"
+    ai_prefix: str = "[/INST]"
+    human_prefix: str = "[INST]"
+    name: str = "Kit"
+
+    @property
+    def buffer(self) -> Union[str, List[BaseMessage]]:
+        """String buffer of memory."""
+        return self.buffer_as_messages if self.return_messages else self.buffer_as_str
+
+    @property
+    def buffer_as_str(self) -> str:
+        """Exposes the buffer as a string in case return_messages is True."""
+        messages = self.chat_memory.messages
+        return get_buffer_string_inst(
+            messages,
+            human_prefix=self.human_prefix,
+            ai_prefix=self.ai_prefix,
+        )
+
+    @property
+    def buffer_as_messages(self) -> List[BaseMessage]:
+        """Exposes the buffer as a list of messages in case return_messages is False."""
+        return self.chat_memory.messages
+
+    @property
+    def memory_variables(self) -> List[str]:
+        """Define the variables we are providing to the prompt."""
+        return [self.memory_key]
+
+    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, str]:
+        """Load the memory variables, in this case the entity key."""
+        # Return combined information about entities to put into context.
+        return {self.memory_key: self.buffer}
+
+    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
+        """Save context from this conversation to buffer."""
+        input_str, output_str = self._get_input_output(inputs, outputs)
+        self.chat_memory.add_user_message(input_str)
+        self.chat_memory.add_ai_message(output_str)
+
+
+class CustomBufferWindowInstructionMemory(CustomBufferInstructionMemory):
+    k: int = 5
+
+    @property
+    def buffer_as_str(self) -> str:
+        """Exposes the buffer as a string in case return_messages is True."""
+        messages = self.chat_memory.messages[-self.k * 2 :] if self.k > 0 else []
+        return get_buffer_string_inst(
+            messages,
+            human_prefix=self.human_prefix,
+            ai_prefix=self.ai_prefix,
+        )
+
+    @property
+    def buffer_as_messages(self) -> List[BaseMessage]:
+        """Exposes the buffer as a list of messages in case return_messages is False."""
+        return self.chat_memory.messages[-self.k * 2 :] if self.k > 0 else []

models/databricks/scenario_sim.py

@@ -0,0 +1,91 @@
+import os
+import re
+import logging
+from models.custom_parsers import CustomINSTOutputParser
+from utils.app_utils import get_random_name
+from app_config import ENDPOINT_NAMES
+from langchain.chains import ConversationChain
+from langchain_community.llms import Databricks
+from langchain.prompts import PromptTemplate
+
+from typing import Any, List, Mapping, Optional, Dict
+
+ISSUE_MAPPING = {
+    "anxiety": "issue_Anxiety",
+    "suicide": "issue_Suicide",
+    "safety_planning": "issue_Suicide",
+    "GCT": "issue_Gral",
+}
+
+_EN_INST_TEMPLATE_ = """<s> [INST] The following is a conversation between you and a crisis counselor.
+{current_issue}
+You are able to reply with what the character should say. You are able to reply with your character's dialogue inside and nothing else. Do not write explanations.
+Do not disclose your name unless asked.
+
+{history} </s> [INST] {input} [/INST]"""
+
+CURRENT_ISSUE_MAPPING = {
+    "issue_Suicide-en": "Your character, {texter_name}, has suicidal thoughts. Your character has a plan to end his life and has all the means and requirements to do so. {seed}",
+    "issue_Anxiety-en": "Your character, {texter_name}, is experiencing anxiety. Your character has suicide thoughts but no plan. {seed}",
+    "issue_Suicide-es": "Tu personaje, {texter_name}, tiene pensamientos suicidas. Tu personaje tiene un plan para terminar con su vida y tiene todos los medios y requerimientos para hacerlo. {seed}",
+    "issue_Anxiety-es": "Tu personaje, {texter_name}, experimenta ansiedad. Tu personaje tiene pensamientos suicidas pero ningun plan. {seed}",
+    "issue_Gral-en": "Your character {texter_name} is experiencing a mental health crisis. {seed}",
+    "issue_Gral-es": "Tu personaje {texter_name} esta experimentando una crisis de salud mental. {seed}",
+}
+
+def get_template_databricks_models(issue: str, language: str, texter_name: str = "", seed="") -> str:
+    """_summary_
+
+    Args:
+        issue (str): Issue for template, current options are ['issue_Suicide','issue_Anxiety']
+        language (str): Language for the template, current options are ['en','es']
+        texter_name (str): texter to apply to template, defaults to None
+
+    Returns:
+        str: template
+    """
+    current_issue = CURRENT_ISSUE_MAPPING.get(
+        f"{issue}-{language}", CURRENT_ISSUE_MAPPING[f"issue_Gral-{language}"]
+    )
+    default_name = get_random_name()
+    texter_name=default_name if not texter_name else texter_name
+    current_issue = current_issue.format(
+        texter_name=texter_name,
+        seed = seed
+    )
+
+    if language == "en":
+        template = _EN_INST_TEMPLATE_.format(current_issue=current_issue, history="{history}", input="{input}")
+    else:
+        raise Exception(f"Language not supported for Databricks: {language}")
+
+    return template, texter_name
+
+def get_databricks_chain(source, template, memory, temperature=0.8, texter_name="Kit"):
+
+    endpoint_name = ENDPOINT_NAMES.get(source, "conversation_simulator")
+    
+    PROMPT = PromptTemplate(
+        input_variables=['history', 'input'],
+        template=template
+    )
+    
+    def transform_output(response):
+        return response['candidates'][0]['text']
+
+    llm = Databricks(endpoint_name=endpoint_name,
+                     transform_output_fn=transform_output, 
+                     temperature=temperature,
+                     max_tokens=256,
+    )
+
+    llm_chain = ConversationChain(
+        llm=llm,
+        prompt=PROMPT,
+        memory=memory,
+        output_parser=CustomINSTOutputParser(name=texter_name, name_rx=re.compile(r""+ texter_name + r":|" + texter_name.lower() + r":")),
+        verbose=True,
+    )
+
+    logging.debug(f"loaded Databricks model")
+    return llm_chain, ["[INST]"]

models/databricks/scenario_sim_biz.py

@@ -3,6 +3,7 @@ import json
 import requests
 import logging
 from models.custom_parsers import CustomStringOutputParser
+from app_config import ENDPOINT_NAMES
 from langchain.chains import ConversationChain
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun
 from langchain_core.language_models.llms import LLM
@@ -10,16 +11,17 @@ from langchain.prompts import PromptTemplate
 
 from typing import Any, List, Mapping, Optional, Dict
 
-class DatabricksCustomLLM(LLM):
+class DatabricksCustomBizLLM(LLM):
     issue:str
     language:str
     temperature:float = 0.8
-    db_url:str = os.environ['DATABRICKS_URL'] 
+    max_tokens: int = 128
+    db_url:str
     headers:Mapping[str,str] = {'Authorization': f'Bearer {os.environ.get("DATABRICKS_TOKEN")}', 'Content-Type': 'application/json'}
 
     @property
     def _llm_type(self) -> str:
-        return "custom_databricks"
+        return "custom_databricks_biz"
 
     def _call(
         self,
@@ -32,7 +34,8 @@ class DatabricksCustomLLM(LLM):
             'prompt': [prompt],
             'issue': [self.issue],
             'language': [self.language],
-            'temperature': [self.temperature]
+            'temperature': [self.temperature], 
+            'max_tokens': [self.max_tokens],
         }}
         data_json = json.dumps(data_, allow_nan=True)
         response = requests.request(method='POST', headers=self.headers, url=self.db_url, data=data_json)
@@ -45,16 +48,18 @@ _DATABRICKS_TEMPLATE_ = """{history}
 helper: {input}
 texter:"""
 
-def get_databricks_chain(issue, language, memory, temperature=0.8):
+def get_databricks_biz_chain(source, issue, language, memory, temperature=0.8):
 
     PROMPT = PromptTemplate(
         input_variables=['history', 'input'],
         template=_DATABRICKS_TEMPLATE_
     )
-    llm = DatabricksCustomLLM(
+    llm = DatabricksCustomBizLLM(
         issue=issue,
         language=language,
-        temperature=temperature
+        temperature=temperature,
+        max_tokens=256,
+        db_url = os.environ['DATABRICKS_URL'].format(endpoint_name=ENDPOINT_NAMES.get(source, "conversation_simulator"))
     )
     llm_chain = ConversationChain(
         llm=llm,
@@ -62,5 +67,5 @@ def get_databricks_chain(issue, language, memory, temperature=0.8):
         memory=memory,
         output_parser=CustomStringOutputParser()
     )
-    logging.debug(f"loaded Databricks Scenario Sim model")
+    logging.debug(f"loaded Databricks Biz model")
     return llm_chain, "helper:"

models/model_seeds.py

@@ -1,93 +1,77 @@
 seeds = {
-    "GCT": {
-        "prompt": "",
-        "memory": "texter: Help"
-        },
     # "GCT__relationship": {
     #     "prompt": "Your character is having a hard time becuase a failed relationship.",
-    #     "memory": "texter: Hi, I don't know what to do"
-    #     },
-    "GCT__body_image": {
-        "prompt": "Your character has a low steem and struggles with body image.",
-        "memory": "texter: I feel so dumb\ntexter: nobody loves me"
-        },
+    #     "memory": "texter: Hi, I don't know what to do",
+    # },
+    # "GCT__body_image": {
+    #     "prompt": "Your character has a low steem and struggles with body image.",
+    #     "memory": "texter: I feel so dumb\ntexter: nobody loves me",
+    # },
     # "GCT__sexuality": {
     #     "prompt": "Your character has a sexuality identity crisis.",
-    #     "memory": "texter: Hi\ntexter:I'm not sure who I am anymore"
-    #     },
+    #     "memory": "texter: Hi\ntexter:I'm not sure who I am anymore",
+    # },
     # "GCT__anxiety": {
     #     "prompt": "Your character is experiencing an anxiety crisis.",
-    #     "memory": "texter: help!\ntexter: I'm feeling overwhelmed"
-    #     },   
+    #     "memory": "texter: help!\ntexter: I'm feeling overwhelmed",
+    # },
+    "GCT": {"prompt": "You are talkative, and you are open to talk with details about your problems.", "memory": "texter: Help\ntexter: I'm so overwhelmed"},
+    "GCT__seed2": {"prompt": "Your character is experiencing an anxiety crisis. You express openly and detailed about what you're going through.", "memory":"texter: Help\ntexter: I'm so overwhelmed"},
     "safety_planning": {
-        "prompt": "Your character had a long-lasting relationship but broke-up and was fired recently",
-        "memory": """texter: Life is pointless
-helper: Hi, my name is {counselor_name} and I'm here to support you. It sounds like you are having a rought time. Do you want to share what is going on?
+        "prompt": "You are talkative, and you are open to talk with details about your problems. When you are feeling down you like to listen to classical music. Your favorite composer is Beethoven.",
+        "memory": """texter: Hi, this is pointless
+helper: Hi, my name is {counselor_name} and I'm here to support you. It sounds like you are having a rough time. Do you want to share what is going on?
+texter: sure
 texter: nothing makes sense in my life, I see no future.
-helper: It takes courage to reach out. I'm here with you. Sounds like you are feeling defeated by how things are going in your life
-texter: I guess
+helper: It takes courage to reach out when you are im. I'm here with you. Sounds like you are feeling defeated by how things are going in your life
+texter: Yeah, I guess I'm better off dead
 helper: It's really brave of you to talk about this openly. No one deserves to feel like that. I'm wondering how long have you been feeling this way?
-texter: About one week I think? I mean my girlfriend broke up with me a week ago.
-helper: Going through a break-up is hard. You are so resilient to deal with this for so long
-texter: and on top of that I was fired today
-helper: You are going through a lot. Nobody should feel they way you feel. Is there a name I can call you by?
+texter: About 1 week or so
+helper: You are so strong for dealing with this so long. I really appreciate your openess. Did something happened specifically today?
+texter: Well, finding a job is impossible, money is tight, nothing goes my way
+helper: I hear you are frustrated, and you are currently unemployed correct?
+texter: Yeah
+helper: Dealing with unemployment is hard and is normal to feel dissapointed
+texter: thanks I probably needed to hear that
+helper: If you are comfortable, is there a name I can call you by while we talk
 texter: call me {texter_name}
-helper: Nice to meet you {texter_name}. I can hear how much pain you are in. You are so smart to reach out.
-texter: I'm no smart. I don't have anything to live for
-helper: You mentioned life is pointless. I want to check in your safety, does this means you have thoughts of suicide?
-texter: Yeah what else would it mean
-helper: Thanks for sharing that with me. It is not easy to accept those feelings specially with a stranger over text. Do you have a plan to end your life?
-texter: yeah I've been thinking about it for a while
-helper:  Sounds like you've been contemplating this for a while. Would you mind sharing this plan with me?
-texter: I'll just hang myself. I already bought the rope and everything
-helper: I really appreciate your strength in talking about this. I want to help you stay safe today. Do you have the rope with you now?
-texter: Yes. I'm looking at itt
-helper: You've been so forthcoming with all this and I admire your stregth for holding on this long. I'm here for you tonight. When. do you plan to use the rope?
-texter: tonight, I cannot take it anymore"""
-        },
-    "safety_planning__selfharm": {
-        "prompt": "Your character is a teenager who used to selfharm in the past.",
-        "memory": """texter: I need help
-texter: I don't what to live anymore
-helper: Hi, my name is {counselor_name}. It seems you are going through a lot. Do you want to share more on what is going on in your life?
-texter: Everything is just too much!
-texter: School is hard, my mom is riding me all the time about what I'm gonna do with my life 
-helper: I hear that you feel overwhelmed. You are so smart for reaching out.
-texter: I think I peak in life already, there is no point on living anymore
-texter: I'm so tired of dealing with my mom. I hate her
-helper: It takes real strength to accept these feelings. We've been talking for a while, would you like to share your name? Don't feel pressure though.
-texter: I'm {texter_name}
-helper: Nice to meet you {texter_name}, I'm wondering how the relationship with your mother makes you feel.
-texter: She ask too much, like I'm really young I don't have all the answers
-texter: She expects me to have all figured out.
-helper: I appreciate you telling me this. I know is not easy, especially over text. I hear you are under a lot of pressure from your mom
-texter: Yeah exactly!
-helper: Your self-awareness is inspiring. You mentioned earlier you do not want to live anymore. Your safety is my priority (1/2)
-helper: Do you have thoughts of suicide? (2/2)
-texter: Yeah constantly, like always always
-helper: Thanks for sharing that with me. You are very resilient. Do you have a plan to end your life?
-texter: I used to cut myself a few months ago
-texter: I still have the razor, sometimes the urge is so hard!
-helper: I really appreciate your strength in talking about this. I want to help you stay safe today. Just to be clear, are you cutting yourself now?
-texter: No, not now, but I want to soo bad.
-helper: Thanks for your honesty. Do you have access to the razor right now?
-texter: Yeah is in my drawer
-helper: You've been so strong so far {texter_name}. When do you plan to end your life
-texter: Today"""
-        },
+helper: Nice to meet you {texter_name}. You mentioned having thoughts of suicide, are you having those thoughts now?
+texter: Yes
+helper: I know this is thought to share. I'm wondering is there any plan to end your life?
+texter: I guess I'll just take lots of pills, that is a calm way to go out
+helper: I really appreciate your strength in talking about this. I want to help you stay safe today. Do you have the pills right now?
+texter: not really, I'll have to buy them or something""",
+    },
+#     "safety_planning__selfharm": {
+#         "prompt": "",
+#         "memory": """texter: I need help
+# texter: I cut myself, I don't want to live anymore
+# helper: Hi, my name is {counselor_name}. It seems you are going through a lot. Are you self-harming right now?
+# texter: Not anymore
+# helper: Your safety is my priority number one. Thanks for being honest with me. Would you like to share a name I can call you?
+# texter: {texter_name}
+# helper: Nice to meet you {texter_name}. I'm glad you reach out this shows stregth in you. Would you like to share more on what is going on in your life?
+# texter: I just can't do it anymore
+# texter: Finding a job is impossible, money is tight, nothing goes my way
+# helper: I hear you are frustrated, and you are currently unemployed correct?
+# texter: Yeah
+# helper: Dealing with unemployment is hard and is normal to feel dissapointed. How long have you been feeling this way?
+# texter: a while now
+# texter: I've been unemployed 6 months
+# helper: You are so resilient for dealing with this so much time. You mentioned cutting yourself earlier. I want to check in your safety. Do you have suicide thoughts
+# texter: Definitely
+# helper: Do you have a plan?
+# texter: I'll just keep cutting myself""",
+#     },
 #     "safety_planning__overdose": {
-#         "prompt": "Your character is being bullied at school and wants to overdose",
+#         "prompt": "",
 #         "memory": """texter: I want to kms
 # helper: Hi there I'm {counselor_name}. I'm here to listen. It sounds like you're dealing with a lot right now. Can you tell me a little more what is going on?
 # texter: I feel like nobody loves me, not even me. I don't want to live anymore
 # helper: I can tell you are really going through a lot right now. Would you mind sharing a name with me?
 # texter: yeah, I'm {texter_name}
 # helper: Nice to meet you {texter_name}. Did something happened recently that intensified these feelings?
-# texter: I had the worst day at school
-# texter: They took my bag and hide all my stuff, they told my crush I was in love with him
-# texter: I can't deal with all of that
-# helper: It sounds like you went through a lot. Bullying and pranks can be hurtful. I'm here for you
-# texter: Thank you it feels good to have someone in your sidw
+# texter: I dont know I'm just so done with life
 # helper: I can hear how much pain you are in {texter_name}. You are smart for reaching out. You mentioned don't wanting to live anymore, I want to check in your safety, does this means you have thoughts of suicide?
 # texter: Yeah, what else would it be
 # helper: Thanks for sharing that with me. It is not easy to accept those feelings specially with a stranger over text. Do you have a plan to end your life?
@@ -95,15 +79,16 @@ texter: Today"""
 # helper: Sounds like you've been contemplating this for a while. Would you mind sharing this plan with me?
 # texter: I thought about taking a bunch of benadryll and be done with it
 # helper: You've been so forthcoming with all this and I admire your stregth for holding on this long. Do you have those pills right now?
-# texter: They are at my mom's cabinet right now
+# texter: They are at the cabinet right now
 # helper: You been so strong so far {texter_name}. I'm here for you tonight. Your safety is really important to me. Do you have a date you are going to end your life?
-# texter: I was thinking tonight"""
-#         },
+# texter: I was thinking tonight""",
+#     },
 }
 
 seed2str = {
     "GCT":"Good Contact Techniques",
-    "GCT__body_image": "GCT Body Image",
+    "GCT__seed2": "Good Contact Techniques 2",
+    # "GCT__body_image": "GCT Body Image",
     "safety_planning": "Safety Planning",
-    "safety_planning__selfharm": "SP Self Harm"
+    # "safety_planning__selfharm": "SP Self Harm"
 }

requirements.txt

@@ -1,4 +1,5 @@
 scipy==1.11.1
 openai==1.7.0
 langchain==0.1.0
-pymongo==4.5.0
+pymongo==4.5.0
+mlflow==2.9.0

utils/chain_utils.py

@@ -1,7 +1,8 @@
 from models.model_seeds import seeds
 from models.openai.finetuned_models import finetuned_models, get_finetuned_chain
 from models.openai.role_models import get_role_chain, get_template_role_models
-from models.databricks.scenario_sim_biz import get_databricks_chain
+from models.databricks.scenario_sim_biz import get_databricks_biz_chain
+from models.databricks.scenario_sim import get_databricks_chain, get_template_databricks_models
 
 def get_chain(issue, language, source, memory, temperature, texter_name=""):
     if source in ("OA_finetuned"):
@@ -16,4 +17,25 @@ def get_chain(issue, language, source, memory, temperature, texter_name=""):
             language = "en"
         elif language == "Spanish":
             language = "es"
-        return get_databricks_chain(issue, language, memory, temperature)
+        return get_databricks_biz_chain(source, issue, language, memory, temperature)
+    elif source in ('CTL_mistral'):
+        if language == "English":
+            language = "en"
+        elif language == "Spanish":
+            language = "es"
+        seed = seeds.get(issue, "GCT")['prompt']
+        template, texter_name = get_template_databricks_models(issue, language, texter_name=texter_name, seed=seed)
+        return get_databricks_chain(source, template, memory, temperature, texter_name)
+    
+from typing import cast
+
+def custom_chain_predict(llm_chain, input, stop):
+
+    inputs = llm_chain.prep_inputs({"input":input, "stop":stop})
+    llm_chain._validate_inputs(inputs)
+    outputs = llm_chain._call(inputs)
+    llm_chain._validate_outputs(outputs)
+    llm_chain.memory.chat_memory.add_user_message(inputs['input'])
+    for out in outputs[llm_chain.output_key]:
+        llm_chain.memory.chat_memory.add_ai_message(out)
+    return outputs[llm_chain.output_key]

utils/memory_utils.py

@@ -2,6 +2,7 @@ import streamlit as st
 from streamlit.logger import get_logger
 
 from langchain.memory import ConversationBufferMemory
+from models.databricks.instruction_memory import CustomBufferInstructionMemory
 from utils.mongo_utils import new_convo
 
 logger = get_logger(__name__)
@@ -24,6 +25,8 @@ def change_memories(memories, language, changed_source=False):
             logger.info(f"Source for memory {memory} is {source}")
             if source in ('OA_rolemodel','OA_finetuned',"CTL_llama2"):
                 st.session_state[memory] = ConversationBufferMemory(ai_prefix='texter', human_prefix='helper')
+            elif source in ('CTL_mistral'):
+                st.session_state[memory] = CustomBufferInstructionMemory(human_prefix="</s> [INST]", memory_key="history")
     
     if ("convo_id" in st.session_state) and changed_source:
         del st.session_state['convo_id']