Update app.py

app.py

@@ -1,45 +1,43 @@
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
 
-# Load Falcon-7B-Instruct (open and free)
-model_name = "tiiuae/falcon-7b-instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load Mistral-7B with 8-bit quantization (saves memory!)
+model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",         # Automatically assign layers to available GPU/CPU
+    load_in_8bit=True,         # Use 8-bit quantization
+    torch_dtype=torch.float16  # Reduce precision to save memory
+)
+
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Define personas
+# Viewpoints
 personas = {
-    "🟢 Optimist": "Someone who sees the good and hopeful side of any issue.",
-    "🔴 Pessimist": "Someone who focuses on the risks, negatives, or problems in any issue.",
-    "🟡 Neutral": "Someone who presents a balanced, unbiased view based on logic."
+    "🟢 Optimist": "Respond with hope and positivity",
+    "🔴 Pessimist": "Respond critically and negatively",
+    "🟡 Neutral": "Respond with a balanced, unbiased tone"
 }
 
-# Prompt template
-def build_prompt(topic, style):
-    return f"""You are a debater. Take this persona: {style}.
-Debate Topic: "{topic}"
-Provide a thoughtful and opinionated response with reasoning.
-Answer:"""
-
-# Generate responses
-def debate(topic):
-    results = {}
-    for label, persona in personas.items():
-        prompt = build_prompt(topic, persona)
-        output = pipe(prompt, max_new_tokens=200, temperature=0.9)[0]["generated_text"]
-        answer = output.split("Answer:")[-1].strip()
-        results[label] = answer
-    return results
+# Generate debate
+def generate_debate(topic):
+    responses = {}
+    for label, style in personas.items():
+        prompt = f"[INST] You are a debater. {style}. Topic: '{topic}'. Give a short opinion. [/INST]"
+        result = pipe(prompt, max_new_tokens=150, temperature=0.7, do_sample=True)[0]["generated_text"]
+        responses[label] = result.strip()
+    return "\n\n".join([f"**{k}**: {v}" for k, v in responses.items()])
 
-# Gradio interface
-def run_debate(topic):
-    responses = debate(topic)
-    return "\n\n".join([f"**{k}**:\n{v}" for k, v in responses.items()])
-
-gr.Interface(
-    fn=run_debate,
+# Gradio app
+demo = gr.Interface(
+    fn=generate_debate,
     inputs=gr.Textbox(label="Enter a Debate Topic"),
     outputs=gr.Markdown(),
-    title="🎙️ Multi-Agent Debate Simulator",
-    description="Simulates multi-perspective debates using Falcon-7B on Hugging Face 🤗."
-).launch()
+    title="🧠 Multi-Agent Debate Simulator (Mistral 7B)",
+    description="Debate with different perspectives using the Mistral-7B-Instruct model (quantized)."
+)
+
+demo.launch()