Update app.py

app.py

@@ -1,31 +1,77 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
+import spaces
 
-# Load the multilingual model for text generation
-model_name = "google/mt5-small"
-text_generator = pipeline("text-generation", model=model_name)
+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 generate_text(prompt, language):
-    # Depending on the language, prepend a language token if needed
-    if language == "Arabic":
-        prompt = "<ar>" + prompt
-    else:  # English
-        prompt = "<en>" + prompt
-    
-    generated_text = text_generator(prompt, max_length=200, num_return_sequences=1)
-    return generated_text[0]['generated_text']
+def load_model(model_name):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map=device,
+        torch_dtype="auto",
+        trust_remote_code=True,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    generator = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        return_full_text=False,
+        max_new_tokens=500,
+        do_sample=True,  # Enable sampling for more creative outputs
+        top_k=50,        # Control diversity
+        top_p=0.95       # Control diversity
+    )
+    return generator
 
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("## Bilingual Text Generator")
-    with gr.Row():
-        language_choice = gr.Radio(choices=["English", "Arabic"], label="Select Language")
-        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="Type your prompt here...")
-        generate_button = gr.Button("Generate Text")
-    
-    output_text = gr.Textbox(label="Generated Text", interactive=False)
-    
-    generate_button.click(generate_text, inputs=[prompt_input, language_choice], outputs=output_text)
+@spaces.GPU
+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_kids_story,
+    inputs=[
+        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=["English", "Arabic"],
+            label="Choose Language",
+            value="English"  # Default to English
+        )
+    ],
+    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=[
+        ["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()