Update app.py

app.py

@@ -169,7 +169,7 @@ print("Running LLM Zephyr")
 llm_zephyr = Llama(model_path=zephyr_model_path,n_gpu_layers=GPU_LAYERS-15,max_new_tokens=256, context_window=4096, n_ctx=4096,n_batch=128,verbose=LLAMA_VERBOSE)
 
 print("Running Yi LLM")
-llm_zephyr = Llama(model_path=yi_model_path,n_gpu_layers=GPU_LAYERS-15,max_new_tokens=256, context_window=4096, n_ctx=4096,n_batch=128,verbose=LLAMA_VERBOSE)
+llm_yi = Llama(model_path=yi_model_path,n_gpu_layers=GPU_LAYERS-15,max_new_tokens=256, context_window=4096, n_ctx=4096,n_batch=128,verbose=LLAMA_VERBOSE)
 
 
 # Mistral formatter
@@ -235,16 +235,18 @@ def generate_local(
         formatted_prompt = format_prompt_zephyr(prompt, history,system_message=sys_message)
         llm = llm_zephyr
     else:
+        sys_message= system_message.replace("##LLM_MODEL###",llm_model).replace("##LLM_MODEL_PROVIDER###",llm_provider)
+        sys_system_understand_message = system_understand_message.replace("##LLM_MODEL###",llm_model).replace("##LLM_MODEL_PROVIDER###",llm_provider)
+        formatted_prompt = format_prompt_mistral(prompt, history,system_message=sys_message,system_understand_message=sys_system_understand_message)
+
         if "yi" in llm_model.lower():
             llm_provider= "01.ai"
             llm_model = "Yi"
+            llm = llm_yi
         else:
             llm_provider= "Mistral"
             llm_model = "Mistral"
-        sys_message= system_message.replace("##LLM_MODEL###",llm_model).replace("##LLM_MODEL_PROVIDER###",llm_provider)
-        sys_system_understand_message = system_understand_message.replace("##LLM_MODEL###",llm_model).replace("##LLM_MODEL_PROVIDER###",llm_provider)
-        formatted_prompt = format_prompt_mistral(prompt, history,system_message=sys_message,system_understand_message=sys_system_understand_message)
-        llm = llm_mistral
+            llm = llm_mistral
         
 
     try: