Update main.py

main.py

@@ -21,6 +21,7 @@ class UserInterface():
             "HuggingFaceH4/zephyr-7b-beta",
             "Open-Orca/Mistral-7B-OpenOrca",
             # "Sharathhebbar24/convo_bot_gpt2",
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
             "meta-llama/Llama-2-7b-chat-hf",
         )
         self.models = st.sidebar.selectbox(
@@ -50,7 +51,7 @@ class UserInterface():
 
         self.model_kwargs = {
             "temperature": self.temperature,
-            "max_length": self.max_token_length
+            "max_new_tokens": self.max_token_length
         }
 
         os.environ['HUGGINGFACEHUB_API_TOKEN'] = os.getenv("HF_KEY")
@@ -87,14 +88,15 @@ class UserInterface():
                  disabled=text_input_visibility
             )
 
+            template = "<|system|>\nYou are a intelligent chatbot and expertise in {context}.</s>\n<|user|>\n{question}.\n<|assistant|>"
 
-            template = """
-            Answer the question based on the context, if you don't know then output "Out of Context"
-            Context: {context}
-            Question: {question}
+            # template = """
+            # Answer the question based on the context, if you don't know then output "Out of Context"
+            # Context: {context}
+            # Question: {question}
 
-            Answer: 
-            """
+            # Answer: 
+            # """
             prompt = PromptTemplate(
                 template=template,
                 input_variables=[