Update app.py

app.py

@@ -3,42 +3,63 @@ import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 
-# Load the model and tokenizer
-model_name = "distilgpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Define the function to generate a response
-def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
+# Function to load model and tokenizer based on selection
+def load_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
+
+# Define the function to generate a response with adjustable parameters and model-specific adjustments
+def generate_response(prompt, model_name, persona="I am a helpful assistant.", temperature=0.7, top_p=0.9, repetition_penalty=1.2, max_length=70):
+    # Load the chosen model and tokenizer
+    tokenizer, model = load_model(model_name)
+    
+    # Adjust the prompt format for DialoGPT
+    if model_name == "microsoft/DialoGPT-small":
+        full_prompt = f"User: {prompt}\nBot:"  # Structure as a conversation
+    else:
+        full_prompt = f"{persona}: {prompt}"  # Standard format for other models
+
+    # Tokenize and generate response
+    inputs = tokenizer(full_prompt, return_tensors="pt")
     outputs = model.generate(
         inputs.input_ids,
-        max_length=70,
+        max_length=max_length,
         do_sample=True,
-        temperature=0.6,
-        top_p=0.9,
-        repetition_penalty=1.2,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
         pad_token_id=tokenizer.eos_token_id
     )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
-    return response
+    
+    # Trim the prompt if it appears in the response
+    if model_name == "microsoft/DialoGPT-small":
+        response_without_prompt = response.split("Bot:", 1)[-1].strip()
+    else:
+        response_without_prompt = response.split(":", 1)[-1].strip()
 
-# Persona-based response function
-def persona_response(prompt, persona="You are a helpful talking dog that answers in short, simple phrases."):
-    full_prompt = f"{persona}: {prompt}"
-    return generate_response(full_prompt)
+    return response_without_prompt if response_without_prompt else "I'm not sure how to respond to that."
 
-# Define Gradio interface function
-def chat_interface(user_input, persona="You are a helpful talking dog that answers in short, simple phrases."):
-    return persona_response(user_input, persona)
+# Define Gradio interface function with model selection
+def chat_interface(user_input, model_choice, persona="I am a helpful assistant", temperature=0.7, top_p=0.9, repetition_penalty=1.2, max_length=50):
+    return generate_response(user_input, model_choice, persona, temperature, top_p, repetition_penalty, max_length)
 
-# Gradio interface setup
+# Set up Gradio interface with model selection and parameter sliders
 interface = gr.Interface(
     fn=chat_interface,
-    inputs=["text", "text"],
+    inputs=[
+        gr.Textbox(label="User Input"),
+        gr.Dropdown(choices=["distilgpt2", "gpt2", "microsoft/DialoGPT-small"], label="Model Choice", value="distilgpt2"),
+        gr.Textbox(label="Persona", value="You are a helpful assistant."),
+        gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.1),
+        gr.Slider(label="Top-p (Nucleus Sampling)", minimum=0.1, maximum=1.0, value=0.9, step=0.1),
+        gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.2, step=0.1),
+        gr.Slider(label="Max Length", minimum=10, maximum=100, value=50, step=5)
+    ],
     outputs="text",
-    title="Simple Chatbot",
-    description="Chat with the bot! Add a persona like 'I am a shopping assistant.'"
+    title="Interactive Chatbot with Model Comparison",
+    description="Chat with the bot! Select a model and adjust parameters to see how they affect the response."
 )
 
 # Launch the Gradio app