Update app.py

app.py

@@ -1,44 +1,37 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load DialoGPT model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
-model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
+# Load GODEL model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("microsoft/GODEL-v1_1-base-seq2seq")
+model = AutoModelForSeq2SeqLM.from_pretrained("microsoft/GODEL-v1_1-base-seq2seq")
 
-# Define chatbot function using OpenAI-style message format
-def chat_with_bot(message, history):
-    # Convert history to a flat string
-    conversation = ""
-    for turn in history:
-        if turn["role"] == "user":
-            conversation += turn["content"] + tokenizer.eos_token
-        elif turn["role"] == "assistant":
-            conversation += turn["content"] + tokenizer.eos_token
+# Define the chatbot function
+def generate_response(message, history):
+    # Format dialog history
+    dialog = [turn["content"] for turn in history if turn["role"] == "user"]
+    dialog.append(message)
+    dialog_text = " EOS ".join(dialog)
 
-    # Append current user message
-    conversation += message + tokenizer.eos_token
+    # GODEL expects an instruction and context
+    instruction = "Instruction: given a dialog context, respond appropriately."
+    query = f"{instruction} [CONTEXT] {dialog_text}"
 
     # Tokenize and generate
-    input_ids = tokenizer.encode(conversation, return_tensors="pt")
+    input_ids = tokenizer(query, return_tensors="pt").input_ids
     output_ids = model.generate(
         input_ids,
-        max_length=1000,
-        pad_token_id=tokenizer.eos_token_id,
+        max_length=128,
         do_sample=True,
-        top_k=50,
-        top_p=0.95
+        top_p=0.9,
+        temperature=0.7
     )
-
-    # Decode only the new response
-    response = tokenizer.decode(output_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
-
+    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return {"role": "assistant", "content": response}
 
 # Launch the chatbot
 gr.ChatInterface(
-    fn=chat_with_bot,
-    title="Muhammad’s ChatGPT",
-    description="A simple chatbot powered by DialoGPT-small. Ask me anything!",
-    type="messages"  # Important for compatibility with message dicts
+    fn=generate_response,
+    title="Muhammad’s GODEL Chatbot",
+    description="A grounded chatbot powered by Microsoft's GODEL model.",
+    type="messages"
 ).launch()