Update app.py

app.py

@@ -1,32 +1,43 @@
 import gradio as gr
+from transformers import pipeline, GPT2LMHeadModel, GPT2Tokenizer
 
 title = "GPT2"
-description = "Gradio Demo for OpenAI GPT2. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."
-
+description = "Gradio Demo for OpenAI GPT2. To use it, simply add your text, or click one of the examples to load them."
 article = "<p style='text-align: center'><a href='https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf' target='_blank'>Language Models are Unsupervised Multitask Learners</a></p>"
 
 examples = [
     ['Paris is the capital of', "gpt2-medium"]
 ]
 
-# Load all models at startup
-io1 = gr.load("huggingface/distilgpt2")
-io2 = gr.load("huggingface/gpt2-large")
-io3 = gr.load("huggingface/gpt2-medium")
-io4 = gr.load("huggingface/gpt2-xl")
+# Initialize models dictionary to cache loaded models
+models = {}
+
+def load_model(model_name):
+    if model_name not in models:
+        tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+        model = GPT2LMHeadModel.from_pretrained(model_name)
+        models[model_name] = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    return models[model_name]
 
-def inference(text, model):
-    if model == "gpt2-large":
-        outtext = io2(text)
-    elif model == "gpt2-medium":
-        outtext = io3(text)
-    elif model == "gpt2-xl":
-        outtext = io4(text)
-    else:
-        outtext = io1(text)
-    return outtext   
+def inference(text, model_name):
+    # Map the model names to their Hugging Face identifiers
+    model_map = {
+        "distilgpt2": "distilgpt2",
+        "gpt2-medium": "gpt2-medium",
+        "gpt2-large": "gpt2-large",
+        "gpt2-xl": "gpt2-xl"
+    }
+    
+    # Get the correct model identifier
+    hf_model_name = model_map.get(model_name, "distilgpt2")
+    
+    # Load the model (will be cached after first load)
+    generator = load_model(hf_model_name)
+    
+    # Generate text
+    generated = generator(text, max_length=50, num_return_sequences=1)
+    return generated[0]['generated_text']
 
-# Create the interface
 iface = gr.Interface(
     inference, 
     [