Update app.py

app.py

@@ -27,15 +27,35 @@ class ConversationManager:
             return self.models[model_name]
         except Exception as e:
             print(f"Failed to load model {model_name}: {e}")
+            print(f"Error type: {type(e).__name__}")
+            print(f"Error details: {str(e)}")
             return None
 
     def generate_response(self, model_name, prompt):
         model, tokenizer = self.load_model(model_name)
-        inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
+    
+        # Format the prompt based on the model
+        if "llama" in model_name.lower():
+            formatted_prompt = self.format_llama2_prompt(prompt)
+        else:
+            formatted_prompt = self.format_general_prompt(prompt)
+    
+        inputs = tokenizer(formatted_prompt, return_tensors="pt", max_length=1024, truncation=True)
         with torch.no_grad():
             outputs = model.generate(**inputs, max_length=200, num_return_sequences=1, do_sample=True)
         return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+    def format_llama2_prompt(self, prompt):
+        B_INST, E_INST = "[INST]", "[/INST]"
+        B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
+        system_prompt = "You are a helpful AI assistant. Please provide a concise and relevant response."
+    
+        formatted_prompt = f"{B_INST} {B_SYS}{system_prompt}{E_SYS}{prompt.strip()} {E_INST}"
+        return formatted_prompt
+
+    def format_general_prompt(self, prompt):
+        # A general format that might work for other models
+        return f"Human: {prompt.strip()}\n\nAssistant:"
     def add_to_conversation(self, model_name, response):
         self.conversation.append((model_name, response))
         if "task complete?" in response.lower():  # Check for task completion marker
@@ -72,48 +92,56 @@ def get_model(dropdown, custom):
     return (model, model)  # Return a tuple (label, value)
 
 def chat(model1, model2, user_input, history, inserted_response=""):
-    model1 = get_model(model1, model1_custom.value)[0]
-    model2 = get_model(model2, model2_custom.value)[0]
-    
-    if not manager.conversation:
-        manager.initial_prompt = user_input
-        manager.clear_conversation()
-        manager.add_to_conversation("User", user_input)
-    
-    models = [model1, model2]
-    current_model_index = 0 if manager.current_model in ["User", "Model 2"] else 1
-    
-    while not manager.task_complete:  # Continue until task is complete
-        if manager.is_paused:
-            yield history, "Conversation paused."
-            return
-
-        model = models[current_model_index]
-        manager.current_model = model
-        
-        if inserted_response and current_model_index == 0:
-            response = inserted_response
-            inserted_response = ""
-        else:
-            prompt = manager.get_conversation_history() + "\n\nPlease continue the conversation. If you believe the task is complete, end your response with 'Task complete?'"
-            response = manager.generate_response(model, prompt)
+    try:
+        model1 = get_model(model1, model1_custom.value)[0]
+        model2 = get_model(model2, model2_custom.value)[0]
         
-        manager.add_to_conversation(model, response)
-        history = manager.get_conversation_history()
+        if not manager.load_model(model1) or not manager.load_model(model2):
+            return "Error: Failed to load one or both models. Please check the model names and try again.", ""
         
-        for i in range(manager.delay, 0, -1):
-            yield history, f"{model} is writing... {i}"
-            time.sleep(1)
+        if not manager.conversation:
+            manager.initial_prompt = user_input
+            manager.clear_conversation()
+            manager.add_to_conversation("User", user_input)
         
-        yield history, ""
+        models = [model1, model2]
+        current_model_index = 0 if manager.current_model in ["User", "Model 2"] else 1
         
-        if manager.task_complete:
-            yield history, "Models believe the task is complete. Are you satisfied with the result? (Yes/No)"
-            return
+        while not manager.task_complete:
+            if manager.is_paused:
+                yield history, "Conversation paused."
+                return
+
+            model = models[current_model_index]
+            manager.current_model = model
+            
+            if inserted_response and current_model_index == 0:
+                response = inserted_response
+                inserted_response = ""
+            else:
+                conversation_history = manager.get_conversation_history()
+                prompt = f"{conversation_history}\n\nPlease continue the conversation. If you believe the task is complete, end your response with 'Task complete?'"
+                response = manager.generate_response(model, prompt)
+            
+            manager.add_to_conversation(model, response)
+            history = manager.get_conversation_history()
+            
+            for i in range(manager.delay, 0, -1):
+                yield history, f"{model} is writing... {i}"
+                time.sleep(1)
+            
+            yield history, ""
+            
+            if manager.task_complete:
+                yield history, "Models believe the task is complete. Are you satisfied with the result? (Yes/No)"
+                return
+            
+            current_model_index = (current_model_index + 1) % 2
         
-        current_model_index = (current_model_index + 1) % 2
-    
-    return history, "Conversation completed."
+        return history, "Conversation completed."
+    except Exception as e:
+        print(f"Error in chat function: {str(e)}")
+        return f"An error occurred: {str(e)}", ""
 
 def user_satisfaction(satisfied, history):
     if satisfied.lower() == 'yes':