Adding options for FlanT5 small, base and large

app.py

@@ -1,29 +1,43 @@
 import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-# Load Flan-T5-base model and tokenizer from Hugging Face
-model_name = "google/flan-t5-base"
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load the shared tokenizer (you can use the tokenizer from any Flan-T5 model)
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
 
+# Define the model names
+model_names = {
+    "Flan-T5-small": "google/flan-t5-small",
+    "Flan-T5-base": "google/flan-t5-base",
+    "Flan-T5-large": "google/flan-t5-large"
+}
+
+# Pre-load the models
+loaded_models = {
+    model_name: AutoModelForSeq2SeqLM.from_pretrained(model_path)
+    for model_name, model_path in model_names.items()
+}
 
 def respond(
     message,
     history: list[tuple[str, str]],
+    model_choice,
     max_tokens,
     temperature,
     top_p,
 ):
+    # Select the pre-loaded model based on user's choice
+    model = loaded_models[model_choice]
+
     # Prepare the input by concatenating the history into a dialogue format
     input_text = ""
     for user_msg, bot_msg in history:
         input_text += f"User: {user_msg} Assistant: {bot_msg} "
     input_text += f"User: {message}"
 
-    # Tokenize the input text
+    # Tokenize the input text using the shared tokenizer
     inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
 
-    # Generate the response using Flan-T5-base
+    # Generate the response using the selected Flan-T5 model
     output_tokens = model.generate(
         inputs["input_ids"],
         max_length=max_tokens,
@@ -36,11 +50,11 @@ def respond(
     response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
     yield response
 
-
 # Define the Gradio interface
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
+        gr.Dropdown(choices=["Flan-T5-small", "Flan-T5-base", "Flan-T5-large"], value="Flan-T5-base", label="Model"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
@@ -51,3 +65,4 @@ if __name__ == "__main__":
     demo.launch()
 
 
+