Update app.py

app.py

@@ -1,60 +1,43 @@
-import gradio as gr
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# ---------------- Agent 1: Intent Classifier ----------------
-intent_classifier = pipeline("zero-shot-classification", model="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli")
-
-
-def detect_intent(text):
-    labels = {
-        "weather": "The user wants to know the weather.",
-        "faq": "The user is asking for help.",
-        "smalltalk": "The user is making casual conversation."
-    }
-    best_intent = "smalltalk"
-    best_score = 0
-    for label, hypothesis in labels.items():
-        result = intent_classifier(text=text, text_pair=hypothesis)[0]
-        if result['label'] == 'ENTAILMENT' and result['score'] > best_score:
-            best_score = result['score']
-            best_intent = label
-    return best_intent
-
-# ---------------- Agent 2: Domain Logic ----------------
-def handle_logic(intent):
-    if intent == "weather":
-        return "It's sunny and 26°C today."
-    elif intent == "faq":
-        return "To reset your password, use the 'Forgot Password' option."
-    else:
-        return "That's great! Anything else you'd like to talk about?"
-
-# ---------------- Agent 3: Natural Language Generation ----------------
-tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
-model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
-
-def generate_reply(prompt):
-    input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors='pt')
-    output_ids = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
-    response = tokenizer.decode(output_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-    return response
-
-# ---------------- Chatbot Pipeline ----------------
-def chatbot(user_input):
-    intent = detect_intent(user_input)
-    logic = handle_logic(intent)
-    response = generate_reply(logic)
-    return response
-
-# ---------------- Gradio UI ----------------
-gr.Interface(
-    fn=chatbot,
-    inputs=gr.Textbox(label="User Input"),
-    outputs=gr.Textbox(label="Chatbot Response"),
-    title="3-Agent Chatbot",
-    description="Intent Detection → Domain Logic → Natural Language Generation"
-).launch()
-
-
-  
+# Load same or different models for each agent
+tokenizer1 = AutoTokenizer.from_pretrained("gpt2")
+model1 = AutoModelForCausalLM.from_pretrained("gpt2")
+
+tokenizer2 = AutoTokenizer.from_pretrained("gpt2-medium")
+model2 = AutoModelForCausalLM.from_pretrained("gpt2-medium")
+
+tokenizer3 = AutoTokenizer.from_pretrained("gpt2-large")
+model3 = AutoModelForCausalLM.from_pretrained("gpt2-large")
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model1, model2, model3 = model1.to(device), model2.to(device), model3.to(device)
+
+def generate_response(model, tokenizer, prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = model.generate(inputs["input_ids"], max_length=100, pad_token_id=tokenizer.eos_token_id)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+import gradio as gr
+
+def multi_agent_chat(user_input):
+    res1 = generate_response(model1, tokenizer1, user_input)
+    res2 = generate_response(model2, tokenizer2, user_input)
+    res3 = generate_response(model3, tokenizer3, user_input)
+    
+    return res1, res2, res3
+
+interface = gr.Interface(
+    fn=multi_agent_chat,
+    inputs=gr.Textbox(lines=2, placeholder="Ask something..."),
+    outputs=[
+        gr.Textbox(label="Agent 1 (GPT-2)"),
+        gr.Textbox(label="Agent 2 (GPT-2 Medium)"),
+        gr.Textbox(label="Agent 3 (GPT-2 Large)")
+    ],
+    title="3-Agent AI Chatbot"
+)
+
+interface.launch()
+