Update app.py

app.py

@@ -23,30 +23,26 @@ def load_model(model_name):
     )
     return generator
 
-# Global storage for models
-models = {
-    "English": load_model("microsoft/Phi-3-mini-4k-instruct"),
-    "Arabic": load_model("ALLaM-AI/ALLaM-7B-Instruct-preview")
-}
 @spaces.GPU
-def generate_story(language, prompt):
-    generator = models[language]
-    output = generator(prompt)
+def generate_text(prompt, model_name):
+    generator = load_model(model_name)
+    messages = [{"role": "user", "content": prompt}]
+    output = generator(messages)
     return output[0]["generated_text"]
 
 # Create Gradio interface
 demo = gr.Interface(
     fn=generate_story,
     inputs=[
-        gr.Radio(choices=["English", "Arabic"], label="Select Language"),
+        gr.Radio(choices=["microsoft/Phi-3-mini-4k-instruct", "ALLaM-AI/ALLaM-7B-Instruct-preview"], label="Select Language"),
         gr.Textbox(lines=2, placeholder="Enter your story prompt here...")
     ],
     outputs=gr.Textbox(label="Generated Story"),
     title="Kids Storyteller",
     description="Choose a language and enter a prompt to generate a fun story for kids!",
     examples=[
-        ["English", "Once upon a time in a magical forest..."],
-        ["Arabic", "في قديم الزمان في غابة سحرية..."]
+        ["microsoft/Phi-3-mini-4k-instruct", "Once upon a time in a magical forest..."],
+        ["ALLaM-AI/ALLaM-7B-Instruct-preview", "في قديم الزمان في غابة سحرية..."]
     ]
 )