Update app.py

app.py

@@ -1,42 +1,62 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+import concurrent.futures
 
-# 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")
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-tokenizer3 = AutoTokenizer.from_pretrained("gpt2-large")
-model3 = AutoModelForCausalLM.from_pretrained("gpt2-large")
+# Load models and tokenizers
+def load_model(name):
+    tokenizer = AutoTokenizer.from_pretrained(name)
+    model = AutoModelForCausalLM.from_pretrained(name)
+    
+    # Define pad token explicitly
+    tokenizer.pad_token = tokenizer.eos_token
+    model.config.pad_token_id = tokenizer.pad_token_id
+    
+    return tokenizer, model.to(device)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model1, model2, model3 = model1.to(device), model2.to(device), model3.to(device)
+tokenizer1, model1 = load_model("gpt2")
+tokenizer2, model2 = load_model("distilgpt2")
+tokenizer3, model3 = load_model("gpt2")
 
+# Generation function
 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)
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
+    outputs = model.generate(
+        inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        max_length=100,
+        pad_token_id=tokenizer.pad_token_id,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9
+    )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-import gradio as gr
-
+# Multi-agent handler
 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
-
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        futures = [
+            executor.submit(generate_response, model1, tokenizer1, user_input),
+            executor.submit(generate_response, model2, tokenizer2, user_input),
+            executor.submit(generate_response, model3, tokenizer3, user_input)
+        ]
+        results = [f.result() for f in futures]
+    return results
+
+# Gradio Interface
 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)")
+        gr.Textbox(label="Agent 2 (DistilGPT-2)"),
+        gr.Textbox(label="Agent 3 (GPT-2)")
     ],
-    title="3-Agent AI Chatbot"
+    title="3-Agent AI Chatbot",
+    description="Three GPT-style agents respond to your input in parallel!"
 )
 
 interface.launch()