app.py

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch

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(
        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

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"]

# Create Gradio interface
demo = gr.Interface(
    fn=generate_text,
    inputs=[
        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="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"]
    ]
)

demo.launch()