Update app.py

app.py

@@ -1,39 +1,43 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, BertForMaskedLM, BertTokenizer
 
-# Load the models with trust_remote_code=True
-model1 = AutoModelForCausalLM.from_pretrained("microsoft/GRIN-MoE", trust_remote_code=True)
-tokenizer1 = AutoTokenizer.from_pretrained("microsoft/GRIN-MoE", trust_remote_code=True)
+# Load GPT-2 for the initial argument
+model1 = GPT2LMHeadModel.from_pretrained("gpt2")
+tokenizer1 = GPT2Tokenizer.from_pretrained("gpt2")
 
-model2 = AutoModelForCausalLM.from_pretrained("microsoft/GRIN-MoE", trust_remote_code=True)
-tokenizer2 = AutoTokenizer.from_pretrained("microsoft/GRIN-MoE", trust_remote_code=True)
+# Load BERT for the counter-argument
+model2 = BertForMaskedLM.from_pretrained("bert-base-uncased")
+tokenizer2 = BertTokenizer.from_pretrained("bert-base-uncased")
 
-def generate_argument(model, tokenizer, prompt, max_length=200, temperature=0.7):
+def generate_gpt2_argument(query, max_length=200, temperature=0.7):
     try:
-        inputs = tokenizer(prompt, return_tensors="pt")
-        outputs = model.generate(
-            **inputs,
+        prompt = f"Provide a logical explanation for the following topic: {query}\n\nExplanation:"
+        inputs = tokenizer1.encode(prompt, return_tensors="pt")
+        outputs = model1.generate(
+            inputs,
             max_length=max_length,
-            num_return_sequences=1,
             temperature=temperature,
-            do_sample=True
+            num_return_sequences=1,
+            do_sample=True,
+            pad_token_id=tokenizer1.eos_token_id
         )
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return tokenizer1.decode(outputs[0], skip_special_tokens=True).replace(prompt, "").strip()
     except Exception as e:
         return f"An error occurred: {str(e)}"
 
-def generate_initial_argument(query):
-    prompt = f"Provide a logical explanation for the following topic: {query}"
-    return generate_argument(model1, tokenizer1, prompt)
-
-def generate_counter_argument(query, initial_argument):
-    prompt = f"Given the topic '{query}' and the initial argument '{initial_argument}', provide a well-reasoned counter-argument:"
-    response = generate_argument(model2, tokenizer2, prompt)
-    return response.replace(prompt, "").strip()
+def generate_bert_argument(query, initial_argument, max_length=200):
+    try:
+        prompt = f"Topic: {query}. Initial argument: {initial_argument}. Counter-argument:"
+        inputs = tokenizer2.encode(prompt, return_tensors="pt", max_length=max_length, truncation=True)
+        outputs = model2(inputs).logits
+        generated_tokens = outputs.argmax(dim=-1)
+        return tokenizer2.decode(generated_tokens[0], skip_special_tokens=True).replace(prompt, "").strip()
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
 
 def debate(query):
-    initial_argument = generate_initial_argument(query)
-    counter_argument = generate_counter_argument(query, initial_argument)
+    initial_argument = generate_gpt2_argument(query)
+    counter_argument = generate_bert_argument(query, initial_argument)
     return initial_argument, counter_argument
 
 # Define the Gradio interface
@@ -41,11 +45,11 @@ iface = gr.Interface(
     fn=debate,
     inputs=gr.Textbox(lines=2, placeholder="Enter your question or topic for debate here..."),
     outputs=[
-        gr.Textbox(label="Initial Argument (GRIN-MoE)"),
-        gr.Textbox(label="Counter-Argument (GRIN-MoE)")
+        gr.Textbox(label="Initial Argument (GPT-2)"),
+        gr.Textbox(label="Counter-Argument (BERT)")
     ],
     title="Two-Model Debate System",
-    description="Enter a question or topic. GRIN-MoE will provide an initial argument and a counter-argument.",
+    description="Enter a question or topic. GPT-2 will provide an initial argument, and BERT will generate a counter-argument.",
     examples=[
         ["What are the long-term implications of artificial intelligence on employment?"],
         ["Should governments prioritize space exploration or addressing climate change?"],