Update app.py

app.py

@@ -3,6 +3,15 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 import spaces
 
+
+def get_model_name(language):
+    """Map language choice to the corresponding model."""
+    model_mapping = {
+        "English": "microsoft/Phi-3-mini-4k-instruct",
+        "Arabic": "ALLaM-AI/ALLaM-7B-Instruct-preview"
+    }
+    return model_mapping.get(language, "ALLaM-AI/ALLaM-7B-Instruct-preview")  # Default to Arabic model
+
 def load_model(model_name):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = AutoModelForCausalLM.from_pretrained(
@@ -18,35 +27,53 @@ def load_model(model_name):
         tokenizer=tokenizer,
         return_full_text=False,
         max_new_tokens=500,
-        do_sample=False
+        do_sample=True,  # Enable sampling for more creative outputs
+        top_k=50,        # Control diversity
+        top_p=0.95       # Control diversity
     )
     return generator
 
+
 @spaces.GPU
-def generate_text(prompt, model_name):
+def generate_kids_story(character, setting, language):
+    model_name = get_model_name(language)
     generator = load_model(model_name)
+
+    # Define prompt for the AI model
+    if language == "English":
+        prompt = (f"Write a short story for kids about a character named {character} who goes on an adventure in {setting}. "
+                  "Make it fun, engaging, and suitable for children.")
+    else:
+        prompt = (f"اكتب قصة قصيرة للأطفال عن شخصية اسمها {character} التي تذهب في مغامرة في {setting}. "
+                  "اجعلها ممتعة وجذابة ومناسبة للأطفال.")
+
     messages = [{"role": "user", "content": prompt}]
     output = generator(messages)
     return output[0]["generated_text"]
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=generate_text,
+    fn=generate_kids_story,
     inputs=[
-        gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
+        gr.Textbox(placeholder="Enter a character name (e.g., Benny the Bunny)...", label="Character Name"),
+        gr.Textbox(placeholder="Enter a setting (e.g., a magical forest)...", label="Setting"),
         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"
+            choices=["English", "Arabic"],
+            label="Choose Language",
+            value="English"  # Default to English
         )
     ],
-    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="Kids' Story"),
+    title="📖 AI Kids' Story Generator - English & Arabic 📖",
+    description="Enter a character name and a setting, and AI will generate a fun short story for kids in English or Arabic.",
     examples=[
-        ["Tell me a funny joke about chickens.", "microsoft/Phi-3-mini-4k-instruct"],
-        ["أخبرني نكتة مضحكة عن الدجاج.", "ALLaM-AI/ALLaM-7B-Instruct-preview"]
-    ]
+        ["Benny the Bunny", "a magical forest", "English"],
+        ["علي البطل", "غابة سحرية", "Arabic"],
+        ["Lila the Ladybug", "a garden full of flowers", "English"],
+        ["ليلى الجنية", "حديقة مليئة بالأزهار", "Arabic"]
+    ],
+    theme="default",
 )
 
+# Launch the Gradio app
 demo.launch()