Update app.py

app.py

@@ -1,44 +1,59 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
-# Load models
-gen_model = pipeline("text2text-generation", model="google/flan-t5-large")
-translator_en_ar = pipeline("translation", model="Helsinki-NLP/opus-mt-en-ar")  # English to Arabic
-translator_ar_en = pipeline("translation", model="Helsinki-NLP/opus-mt-ar-en")  # Arabic to English
-tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
+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 get_plant_info(plant_name, language):
-    if language == "English":
-        prompt = (
-            f"Provide detailed information about {plant_name}. "
-            f"Include its scientific name, growing conditions (light, water, soil type), "
-            f"common uses, and how to take care of it."
-        )
-        response = gen_model(prompt, min_length=50, max_length=300)[0]["generated_text"]
-    else:  # Arabic
-        translated_name = translator_ar_en(plant_name)[0]["translation_text"]  # Convert Arabic input to English
-        prompt = (
-            f"Provide detailed information about {translated_name}. "
-            f"Include its scientific name, growing conditions (light, water, soil type), "
-            f"common uses, and how to take care of it."
-        )
-        response_en = gen_model(prompt, min_length=50, max_length=300)[0]["generated_text"]
-        response = translator_en_ar(response_en)[0]["translation_text"]  # Convert English output back to Arabic
+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=False
+    )
+    return generator
 
-    return response
+def generate_text(prompt, language):
+    model_name = get_model_name(language)  # Get the model based on language choice
+    generator = load_model(model_name)
+    messages = [{"role": "user", "content": prompt}]
+    output = generator(messages)
+    return output[0]["generated_text"]
 
-# Gradio UI
-interface = gr.Interface(
-    fn=get_plant_info,
+# Create Gradio interface
+demo = gr.Interface(
+    fn=generate_text,
     inputs=[
-        gr.Textbox(label="Enter Plant Name / أدخل اسم النبات"),
-        gr.Radio(["English", "العربية"], label="Choose Language / اختر اللغة")
+        gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
+        gr.Dropdown(
+            choices=["English", "Arabic"],  # Users choose the language, not the model name
+            label="Choose Language",
+            value="Arabic"  # Default to Arabic (ALLAM model)
+        )
     ],
-    outputs=gr.Textbox(label="Plant Information / معلومات النبات", lines=10),
-    title="Plant Information App",
-    description="Enter a plant name and select a language to get detailed information."
+    outputs=gr.Textbox(label="Generated Text"),
+    title="Text Generator",
+    description="Enter a prompt and generate text in English or Arabic using AI models.",
+    examples=[
+        ["Give me information about Lavender.", "English"],
+        ["أعطني معلومات عن نبات اللافندر", "Arabic"]
+    ]
 )
 
-# Launch the app
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()