Adding model selector for small, medium and large

app.py

@@ -6,12 +6,15 @@ from langchain.memory import ConversationBufferMemory
 # Move model to device (GPU if available)
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
-# Load the tokenizer (you can use the pre-trained tokenizer for GPT-2 family)
-tokenizer = GPT2Tokenizer.from_pretrained("microsoft/DialoGPT-medium")
-model = GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-medium")
+# Load all three DialoGPT models (small, medium, large)
+models = {
+    "small": GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-small").to(device),
+    "medium": GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-medium").to(device),
+    "large": GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-large").to(device)
+}
 
-# Move model to the device (GPU or CPU)
-model.to(device)
+# Load the tokenizer (same tokenizer for all models)
+tokenizer = GPT2Tokenizer.from_pretrained("microsoft/DialoGPT-medium")
 
 # Set up conversational memory using LangChain's ConversationBufferMemory
 memory = ConversationBufferMemory()
@@ -28,7 +31,7 @@ def truncate_history_to_100_tokens(history, tokenizer, max_tokens=100):
     return tokenized_history
 
 # Define the chatbot function with memory and additional parameters
-def chat_with_dialogpt(input_text, temperature, top_p, top_k):
+def chat_with_dialogpt(input_text, temperature, top_p, top_k, model_size):
     # Retrieve conversation history
     conversation_history = memory.load_memory_variables({})['history']
     
@@ -45,6 +48,9 @@ def chat_with_dialogpt(input_text, temperature, top_p, top_k):
     # Concatenate truncated history with the new input
     final_input_ids = torch.cat((truncated_input_ids_tensor, input_ids), dim=1)
 
+    # Get the model corresponding to the selected size
+    model = models[model_size]
+
     # Generate the response using the model with adjusted parameters
     outputs = model.generate(
         final_input_ids,
@@ -83,6 +89,9 @@ with gr.Blocks() as interface:
     
     # Add the instruction message above the input box
     gr.Markdown("**Instructions:** Press `Shift + Enter` to submit, and `Enter` for a new line.")
+    
+    # Add a dropdown for selecting the model size (small, medium, large)
+    model_selector = gr.Dropdown(choices=["small", "medium", "large"], value="medium", label="Select Model Size")
 
     # Add a clear history button
     clear_button = gr.Button("Clear History", scale=0)
@@ -97,17 +106,20 @@ with gr.Blocks() as interface:
     top_k_slider = gr.Slider(1, 100, step=1, value=50, label="Top-k", scale=0)
     
     # Define the function to update the chat
-    def update_chat(input_text, chat_history, temperature, top_p, top_k):
-        updated_history = chat_with_dialogpt(input_text, temperature, top_p, top_k)
+    def update_chat(input_text, chat_history, temperature, top_p, top_k, model_size):
+        updated_history = chat_with_dialogpt(input_text, temperature, top_p, top_k, model_size)
         return updated_history, ""
     
     # Submit when pressing Shift + Enter
     user_input.submit(update_chat, 
-                      inputs=[user_input, chatbot_output, temperature_slider, top_p_slider, top_k_slider], 
+                      inputs=[user_input, chatbot_output, temperature_slider, top_p_slider, top_k_slider, model_selector], 
                       outputs=[chatbot_output, user_input])
 
     # Layout for sliders and chatbot UI
     gr.Row([temperature_slider, top_p_slider, top_k_slider])
-    
+
+    # Layout for model selector and clear button in a row
+    gr.Row([model_selector, clear_button])
+
 # Launch the Gradio app
 interface.launch()