Adding ability to switch between small, base and large models

app.py

@@ -1,16 +1,25 @@
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 from langchain.memory import ConversationBufferMemory
+import torch
 
-# Load the tokenizer and model for flan-t5
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load all three Flan-T5 models (small, base, large)
+models = {
+    "small": T5ForConditionalGeneration.from_pretrained("google/flan-t5-small").to(device),
+    "base": T5ForConditionalGeneration.from_pretrained("google/flan-t5-base").to(device),
+    "large": T5ForConditionalGeneration.from_pretrained("google/flan-t5-large").to(device)
+}
+
+# Load the tokenizer (same tokenizer for all models)
 tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-base")
-model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-base")
 
 # Set up conversational memory using LangChain's ConversationBufferMemory
 memory = ConversationBufferMemory()
 
-# Define the chatbot function with memory
-def chat_with_flan(input_text):
+# Define the chatbot function with memory and model size selection
+def chat_with_flan(input_text, model_size):
     # Retrieve conversation history and append the current user input
     conversation_history = memory.load_memory_variables({})['history']
     
@@ -20,6 +29,9 @@ def chat_with_flan(input_text):
     # Tokenize the input for the model
     input_ids = tokenizer.encode(full_input, return_tensors="pt")
     
+    # Get the model based on the selected size
+    model = models[model_size]
+    
     # Generate the response from the model
     outputs = model.generate(input_ids, max_length=200, num_return_sequences=1)
     
@@ -38,17 +50,22 @@ 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, base, large)
+    model_selector = gr.Dropdown(choices=["small", "base", "large"], value="base", label="Select Model Size")
+
     # Input box for the user
     user_input = gr.Textbox(label="Your Input", placeholder="Type your message here...", lines=2, show_label=True)
     
-    def update_chat(input_text, chat_history):
-        updated_history = chat_with_flan(input_text)
+    # Define the function to update the chat based on selected model
+    def update_chat(input_text, model_size):
+        updated_history = chat_with_flan(input_text, model_size)
         return updated_history, ""
 
     # Submit when pressing Enter
-    user_input.submit(update_chat, inputs=[user_input, chatbot_output], outputs=[chatbot_output, user_input])
+    user_input.submit(update_chat, inputs=[user_input, model_selector], outputs=[chatbot_output, user_input])
 
+    # Layout for model selector and chatbot UI
+    gr.Row([model_selector])
+    
 # Launch the Gradio app
 interface.launch()
-
-