Update app.py

app.py

@@ -18,6 +18,50 @@ models=[
 ]
 
 
+def format_prompt_default(message, history,cust_p):
+    prompt = ""
+    if history:
+        #<start_of_turn>userHow does the brain work?<end_of_turn><start_of_turn>model
+        for user_prompt, bot_response in history:
+            prompt += f"{user_prompt}\n"
+            print(prompt)
+            prompt += f"{bot_response}\n"
+            print(prompt)
+    #prompt += f"{message}\n"
+    prompt+=cust_p.replace("USER_INPUT",message)
+    return prompt
+
+def format_prompt_gemma(message, history,cust_p):
+    prompt = ""
+    if history:
+        #<start_of_turn>userHow does the brain work?<end_of_turn><start_of_turn>model
+        for user_prompt, bot_response in history:
+            prompt += f"{user_prompt}\n"
+            print(prompt)
+            prompt += f"{bot_response}\n"
+            print(prompt)
+    #prompt += f"<start_of_turn>user{message}<end_of_turn><start_of_turn>model"
+    prompt+=cust_p.replace("USER_INPUT",message)
+    return prompt
+
+   
+def format_prompt_mixtral(message, history,cust_p):
+    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]"
+    prompt+=cust_p.replace("USER_INPUT",message)    
+    return prompt
+
+def format_prompt_choose(message, history, cust_p, model_name):
+    if "gemma" in models[model_name].lower():
+        return format_prompt_gemma(message,history,cust_p)
+    if "mixtral" in models[model_name].lower():
+        return format_prompt_mixtral(message,history,cust_p)
+    else:
+        return format_prompt_default(message,history,cust_p)
 
 def load_models(inp):
     print(type(inp))
@@ -25,7 +69,14 @@ def load_models(inp):
     print(models[inp])
     model_state= InferenceClient(models[inp])
     out_box=gr.update(label=models[inp])
-    return out_box, model_state
+    if "gemma" in models[inp].lower():
+        prompt_out="<start_of_turn>userUSER_INPUT<end_of_turn><start_of_turn>model"
+    if "mixtral" in models[inp].lower():
+        prompt_out="[INST] USER_INPUT [/INST]"
+    else:
+        prompt_out="USER_INPUT\n"
+    
+    return out_box,prompt_out, model_state
 
 VERBOSE=False
 
@@ -50,8 +101,9 @@ def format_prompt(message, history, cust_p):
     prompt+=cust_p.replace("USER_INPUT",message)
     return prompt
 
-def chat_inf(system_prompt,prompt,history,memory,model_state,seed,temp,tokens,top_p,rep_p,chat_mem,cust_p):
+def chat_inf(system_prompt,prompt,history,memory,model_state,model_name,seed,temp,tokens,top_p,rep_p,chat_mem,cust_p):
     #token max=8192
+    model_n=models[model_name]
     print(model_state)
     hist_len=0
     client=model_state
@@ -79,9 +131,9 @@ def chat_inf(system_prompt,prompt,history,memory,model_state,seed,temp,tokens,to
             seed=seed,
         )
         if system_prompt:
-            formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", memory[0-chat_mem:],cust_p)
+            formatted_prompt = format_prompt_choose(f"{system_prompt}, {prompt}", memory[0-chat_mem:],cust_p,model_n)
         else:
-            formatted_prompt = format_prompt(prompt, memory[0-chat_mem:],cust_p)
+            formatted_prompt = format_prompt_choose(prompt, memory[0-chat_mem:],cust_p,model_n)
         stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
         output = ""
         for response in stream:
@@ -158,13 +210,13 @@ with gr.Blocks() as app:
                     chatblock=gr.Dropdown(label="Chatblocks",info="Choose specific blocks of chat",choices=[c for c in range(1,40)],multiselect=True)
 
     
-    client_choice.change(load_models,client_choice,[chat_b,model_state])
-    app.load(load_models,client_choice,[chat_b,model_state])
+    client_choice.change(load_models,client_choice,[chat_b,custom_prompt,model_state])
+    app.load(load_models,client_choice,[chat_b,custom_prompt,model_state])
     
     im_go=im_btn.click(get_screenshot,[chat_b,im_height,im_width,chatblock,theme,wait_time],img)
     
-    chat_sub=inp.submit(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,memory,model_state,seed,temp,tokens,top_p,rep_p,chat_mem,custom_prompt],[chat_b,memory])
-    go=btn.click(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,memory,model_state,seed,temp,tokens,top_p,rep_p,chat_mem,custom_prompt],[chat_b,memory])
+    chat_sub=inp.submit(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,memory,model_state,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem,custom_prompt],[chat_b,memory])
+    go=btn.click(check_rand,[rand,seed],seed).then(chat_inf,[sys_inp,inp,chat_b,memory,model_state,client_choice,seed,temp,tokens,top_p,rep_p,chat_mem,custom_prompt],[chat_b,memory])
     
     stop_btn.click(None,None,None,cancels=[go,im_go,chat_sub])
     clear_btn.click(clear_fn,None,[inp,sys_inp,chat_b,memory])