Update app.py

app.py

@@ -39,7 +39,7 @@ def load_models(inp):
     return out_box[0],out_box[1],out_box[2],out_box[3]
 
 
-def format_prompt(message, history):
+def format_prompt_gemma(message, history):
     prompt = ""
     if history:
         #<start_of_turn>userHow does the brain work?<end_of_turn><start_of_turn>model
@@ -49,9 +49,27 @@ def format_prompt(message, history):
             prompt += f"{bot_response}\n"
             print(prompt)
     prompt += f"<start_of_turn>user{message}<end_of_turn><start_of_turn>model"
-    print(prompt)
     return prompt
 
+   
+def format_prompt_mixtral(message, history):
+    prompt = "<s>"
+    if history:
+    
+    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
+def format_prompt_choose(message, history, model_name):
+    if "gemma" in model_name.lower():
+        return format_prompt_gemma(message,history)
+    if "mixtral" in model_name.lower():
+        return format_prompt_mixtral(message,history)
+    else:
+        return format_prompt_mixtral(message,history)
+
+
 
 mega_hist=[[],[],[],[]]
 def chat_inf_tree(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p,rep_p,hid_val):
@@ -99,7 +117,7 @@ def chat_inf_a(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p
             seed=seed,
         )
         #formatted_prompt=prompt   
-        formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+        formatted_prompt = format_prompt_choose(f"{system_prompt}, {prompt}", history, client_choice[0])
         stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
         output = ""
         for response in stream:
@@ -126,7 +144,7 @@ def chat_inf_b(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p
             seed=seed,
         )
         #formatted_prompt=prompt   
-        formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+        formatted_prompt = format_prompt_choose(f"{system_prompt}, {prompt}", history, client_choice[1])
         stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
         output = ""
         for response in stream:
@@ -152,7 +170,7 @@ def chat_inf_c(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p
             seed=seed,
         )
         #formatted_prompt=prompt   
-        formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+        formatted_prompt = format_prompt_choose(f"{system_prompt}, {prompt}", history, client_choice[2])
         stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
         output = ""
         for response in stream:
@@ -178,7 +196,7 @@ def chat_inf_d(system_prompt,prompt,history,client_choice,seed,temp,tokens,top_p
             seed=seed,
         )
         #formatted_prompt=prompt   
-        formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+        formatted_prompt = format_prompt_choose(f"{system_prompt}, {prompt}", history, client_choice[3])
         stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
         output = ""
         for response in stream: