Update app.py

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
-import spaces
 
+# Load models for English and Arabic
 def load_model(model_name):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = AutoModelForCausalLM.from_pretrained(
@@ -18,34 +18,34 @@ def load_model(model_name):
         tokenizer=tokenizer,
         return_full_text=False,
         max_new_tokens=500,
-        do_sample=False
+        do_sample=True  # Changed to enable sampling for more creative outputs
     )
     return generator
 
-@spaces.GPU
-def generate_text(prompt, model_name):
-    generator = load_model(model_name)
-    messages = [{"role": "user", "content": prompt}]
-    output = generator(messages)
+# Global storage for models
+models = {
+    "English": load_model("microsoft/Phi-3-mini-4k-instruct"),
+    "Arabic": load_model("ALLaM-AI/ALLaM-7B-Instruct-preview")
+}
+
+def generate_story(language, prompt):
+    generator = models[language]
+    output = generator(prompt)
     return output[0]["generated_text"]
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=generate_text,
+    fn=generate_story,
     inputs=[
-        gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
-        gr.Dropdown(
-            choices=["Qwen/Qwen2.5-1.5B-Instruct","microsoft/Phi-3-mini-4k-instruct", "ALLaM-AI/ALLaM-7B-Instruct-preview"],
-            label="Choose Model",
-            value="ALLaM-AI/ALLaM-7B-Instruct-preview"
-        )
+        gr.Radio(choices=["English", "Arabic"], label="Select Language"),
+        gr.Textbox(lines=2, placeholder="Enter your story prompt here...")
     ],
-    outputs=gr.Textbox(label="Generated Text"),
-    title="Text Generator",
-    description="Enter a prompt and generate text using one of the available models.",
+    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=[
-        ["Tell me a funny joke about chickens.", "microsoft/Phi-3-mini-4k-instruct"],
-        ["أخبرني نكتة مضحكة عن الدجاج.", "ALLaM-AI/ALLaM-7B-Instruct-preview"]
+        ["English", "Once upon a time in a magical forest..."],
+        ["Arabic", "في قديم الزمان في غابة سحرية..."]
     ]
 )