Update app.py

app.py

@@ -1,37 +1,40 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Add a print statement to confirm initialization
-print("Initializing model...")
-
-# Load AI model
+# Initialize the model and tokenizer
 try:
-    model = pipeline("text2text-generation", model="google/flan-t5-base")
+    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 loading model: {e}")
+    print(f"Error during model initialization: {e}")
+    chat_pipeline = None
 
 def ai_vote(poll_title, choices, num_ais):
-    try:
-        results = {}
-        explanations = []
+    if chat_pipeline is None:
+        return {"Error": "Model not initialized"}, []
 
-        for _ in range(num_ais):
-            input_text = f"Poll Title: {poll_title}\nChoices: {', '.join(choices)}\nChoose the best option and explain why."
-            response = model(input_text, max_length=100, num_return_sequences=1)[0]['generated_text']
-            
-            # Extract the chosen option and explanation
-            chosen_option = response.split("\n")[0].strip()
-            explanation = "\n".join(response.split("\n")[1:]).strip()
+    results = {}
+    explanations = []
 
-            results[chosen_option] = results.get(chosen_option, 0) + 1
-            explanations.append((chosen_option, explanation))
+    for _ 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']
+        except Exception as e:
+            return {"Error": str(e)}, []
 
-        return results, explanations
-    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
 
-# Gradio interface
 def gradio_interface(title, choices, num_ais):
     try:
         choices = [choice.strip() for choice in choices.split(",")]
@@ -48,12 +51,11 @@ interface = gr.Interface(
         gr.Slider(label="Number of AIs", minimum=1, maximum=10, step=1)
     ],
     outputs=[
-        gr.Label(label="Poll Results"),
+        gr.JSON(label="Poll Results"),
         gr.Textbox(label="AI Explanations")
     ]
 )
 
-# Launch app
 if __name__ == "__main__":
     print("Launching interface...")
     interface.launch()