Update app.py

app.py

@@ -3,36 +3,36 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # Initialize the model and tokenizer
 try:
-    print("Initializing model...")
     tokenizer = AutoTokenizer.from_pretrained("satvikag/chatbot")
     model = AutoModelForCausalLM.from_pretrained("satvikag/chatbot")
     chat_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
-    print("Model loaded successfully.")
 except Exception as e:
-    print(f"Error during model initialization: {e}")
+    print(f"Error initializing model: {e}")
     chat_pipeline = None
 
 def ai_vote(poll_title, choices, num_ais):
     if chat_pipeline is None:
         return {"Error": "Model not initialized"}, []
 
-    results = {}
+    # Initialize results and explanations
+    results = {choice: 0 for choice in choices}
     explanations = []
 
-    for _ in range(num_ais):
+    # Loop through the number of AIs to get responses
+    for i in range(num_ais):
         input_text = f"Poll Title: {poll_title}\nChoices: {', '.join(choices)}\nChoose the best option and explain why."
         try:
             response = chat_pipeline(input_text, max_length=150, num_return_sequences=1)[0]['generated_text']
+            # Extract the choice and explanation
+            for choice in choices:
+                if choice.lower() in response.lower():
+                    results[choice] += 1
+                    explanation = response.split("\n", 1)[-1].strip()  # Extract explanation
+                    explanations.append((choice, explanation))
+                    break
         except Exception as e:
             return {"Error": str(e)}, []
 
-        # Extract the chosen option and explanation
-        chosen_option = response.split("\n")[0].strip()
-        explanation = "\n".join(response.split("\n")[1:]).strip()
-
-        results[chosen_option] = results.get(chosen_option, 0) + 1
-        explanations.append((chosen_option, explanation))
-
     return results, explanations
 
 def gradio_interface(title, choices, num_ais):
@@ -43,6 +43,7 @@ def gradio_interface(title, choices, num_ais):
     except Exception as e:
         return {"Error": str(e)}, "An error occurred."
 
+# Gradio Interface
 interface = gr.Interface(
     fn=gradio_interface,
     inputs=[
@@ -57,5 +58,4 @@ interface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    print("Launching interface...")
     interface.launch()