edit to handle mistral prompt format

app.py

@@ -5,6 +5,7 @@ from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
 from langchain.llms import HuggingFaceHub
 from langchain.prompts.chat import (
+PromptTemplate,
 ChatPromptTemplate,
 MessagesPlaceholder,
 SystemMessagePromptTemplate,
@@ -16,6 +17,9 @@ from langchain.schema import AIMessage, HumanMessage
 from openai.error import AuthenticationError
 import streamlit as st
 
+from langchain import verbose
+verbose = True
+
 
 def setup_memory():
     msgs = StreamlitChatMessageHistory(key="basic_chat_app")
@@ -26,42 +30,55 @@ def setup_memory():
     return memory
     
 
-def use_existing_chain(model, provider, temp):
+def use_existing_chain(model, provider, temp, max_tokens):
+    # TODO: consider whether prompt needs to be checked here
+    if "mistral" in model:
+        return False
     if "current_chain" in st.session_state:
         current_chain = st.session_state.current_chain
         if (current_chain.model == model) \
             and (current_chain.provider == provider) \
-            and (current_chain.temp == temp):
+            and (current_chain.temp == temp) \
+            and (current_chain.max_tokens == max_tokens):
             return True
     return False
 
 
 class CurrentChain():
-    def __init__(self, model, provider, memory, temp):
+    def __init__(self, model, provider, prompt, memory, temp, max_tokens=64):
         self.model = model
         self.provider = provider
         self.temp = temp
+        self.max_tokens=max_tokens
+   
         logging.info(f"setting up new chain with params {model_name}, {provider}, {temp}")
         if provider == "OpenAI":                               
             llm = ChatOpenAI(model_name=model, temperature=temp)
         elif provider == "HuggingFace":
+            # TODO: expose the controls below as widgets and clean up init
             llm = HuggingFaceHub(repo_id=model,
-                                 model_kwargs={"temperature": temp, "max_length": 64})
-        prompt = ChatPromptTemplate(
-        messages=[
-            SystemMessagePromptTemplate.from_template(
-                "You are a nice chatbot having a conversation with a human."
-            ),
-            MessagesPlaceholder(variable_name="chat_history"),
-            HumanMessagePromptTemplate.from_template("{input}")
-        ]
-        )
+                                 model_kwargs={"temperature": temp, 
+                                               "max_new_tokens": 256,
+                                              "top_p" : 0.95,
+                                               "repetition_penalty" : 1.0,
+                                              "do_sample" : True,
+                                              "seed" : 42})
+
         self.conversation = LLMChain(
                         llm=llm,
                         prompt=prompt,
                         verbose=True,
                         memory=memory
                     )
+
+
+def format_mistral_prompt(message, history):
+  prompt = "<s>"
+  for user_prompt, bot_response in history:
+    prompt += f"[INST] {user_prompt} [/INST]"
+    prompt += f" {bot_response}</s> "
+  prompt += f"[INST] {message} [/INST]"
+  return prompt
         
   
 if __name__ == "__main__":
@@ -79,7 +96,10 @@ if __name__ == "__main__":
     model_name = st.sidebar.selectbox(
         label = "Choose a model",
         options = ["gpt-3.5-turbo (OpenAI)", 
-                   "bigscience/bloom (HuggingFace)"
+                   # "bigscience/bloom (HuggingFace)",  # runs
+                   # "microsoft/DialoGPT-medium (HuggingFace)",   # throws error
+                   # "google/flan-t5-xxl (HuggingFace)",  # runs
+                   "mistralai/Mistral-7B-Instruct-v0.1 (HuggingFace)"
                   ],
         help="Which LLM to use",
     )
@@ -92,39 +112,87 @@ if __name__ == "__main__":
         min_value=float(0),
         max_value=1.0,
         step=0.1,
-        value=0.9,
+        value=0.4,
         help="Set the decoding temperature"
         )
+
+    max_tokens = st.sidebar.slider(
+        label="Max tokens",
+        min_value=32,
+        max_value=2048,
+        step=1,
+        value=1028,
+        help="Set the maximum number of tokens to generate"
+        )  # TODO: edit this, not currently using
     ##########################
+    
     model = model_name.split("(")[0].rstrip()   # remove name of model provider
     provider = model_name.split("(")[-1].split(")")[0]
+
     if "session_memory" not in st.session_state:
-        st.session_state.session_memory = setup_memory()
-            
-    if use_existing_chain(model, provider, temp):
+        st.session_state.session_memory = setup_memory()  # for openai
+
+    if "history" not in st.session_state:
+        st.session_state.history = []   # for mistral
+
+    if "mistral" in model:
+        prompt = PromptTemplate(input_variables=["input"], 
+                                        template="{input}")
+    else:
+        prompt = ChatPromptTemplate(
+            messages=[
+                SystemMessagePromptTemplate.from_template(
+                    "You are a nice chatbot having a conversation with a human."
+                ),
+                MessagesPlaceholder(variable_name="chat_history"),
+                HumanMessagePromptTemplate.from_template("{input}")
+            ],
+                verbose=True
+            )  
+        
+    if use_existing_chain(model, provider, temp, max_tokens):
         chain = st.session_state.current_chain
     else:
         chain =  CurrentChain(model, 
                      provider, 
+                    prompt,
                      st.session_state.session_memory, 
-                     temp)
+                     temp,
+                    max_tokens)
         st.session_state.current_chain = chain 
 
     conversation = chain.conversation
+        
     if st.button("Clear history"):
-        conversation.memory.clear()
-    for message in conversation.memory.buffer:    # display chat history
-        st.chat_message(message.type).write(message.content)
+        conversation.memory.clear()   # for openai
+        st.session_state.history = []  # for mistral
+        logging.info("history cleared")
+        
+    for user_msg, asst_msg in st.session_state.history:
+        with st.chat_message("user"):
+            st.write(user_msg)
+        with st.chat_message("assistant"):
+            st.write(asst_msg)
+        
     text = st.chat_input()
     if text:
         with st.chat_message("user"):
             st.write(text)
+            logging.info(text)
         try:
-            result = conversation.predict(input=text)
+            if "mistral" in model:
+                full_prompt = format_mistral_prompt(text, st.session_state.history)
+                result = conversation.predict(input=full_prompt)
+            else:
+                result = conversation.predict(input=text)
+
+            st.session_state.history.append((text, result))
+            logging.info(repr(result))
             with st.chat_message("assistant"):
                 st.write(result)
         except (AuthenticationError, ValueError):
-            st.warning("Enter a valid API key", icon="⚠️")
+            st.warning("Supply a valid API key", icon="⚠️")
+