Only loading one model at a time, adding Large and XL models

app.py

@@ -1,20 +1,32 @@
 import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-# Load the shared tokenizer (you can use the tokenizer from any Flan-T5 model)
+# Load the shared tokenizer (can be reused across all models)
 tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
 
-# Define the model names
+# Define the available model names and paths
 model_names = {
     "Flan-T5-small": "google/flan-t5-small",
-    "Flan-T5-base": "google/flan-t5-base"
+    "Flan-T5-base": "google/flan-t5-base",
+    "Flan-T5-large": "google/flan-t5-large",
+    "Flan-T5-XL": "google/flan-t5-xl"
 }
 
-# Pre-load the models
-loaded_models = {
-    model_name: AutoModelForSeq2SeqLM.from_pretrained(model_path)
-    for model_name, model_path in model_names.items()
-}
+# Initialize variables to manage loaded model
+current_model = None
+current_model_name = None
+
+def load_model(model_name):
+    """Load the model if not already loaded or if switching models."""
+    global current_model, current_model_name
+
+    # Load the model only if it hasn't been loaded or if a different one is selected
+    if model_name != current_model_name:
+        print(f"Loading {model_name}...")
+        current_model = AutoModelForSeq2SeqLM.from_pretrained(model_names[model_name])
+        current_model_name = model_name
+
+    return current_model
 
 def respond(
     message,
@@ -24,8 +36,8 @@ def respond(
     temperature,
     top_p,
 ):
-    # Select the pre-loaded model based on user's choice
-    model = loaded_models[model_choice]
+    # Load the selected model (or switch models if needed)
+    model = load_model(model_choice)
 
     # Prepare the input by concatenating the history into a dialogue format
     input_text = ""
@@ -53,7 +65,11 @@ def respond(
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Dropdown(choices=["Flan-T5-small", "Flan-T5-base"], value="Flan-T5-base", label="Model"),
+        gr.Dropdown(
+            choices=["Flan-T5-small", "Flan-T5-base", "Flan-T5-large", "Flan-T5-XL"],
+            value="Flan-T5-base",  # Default selection
+            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)"),
@@ -62,6 +78,3 @@ demo = gr.ChatInterface(
 
 if __name__ == "__main__":
     demo.launch()
-
-
-