app.py

import gradio as gr
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
import os
import torch
import spaces

# Define model paths
MODEL_PATHS = {
    "Terjman-Nano-v2": "BounharAbdelaziz/Terjman-Nano-v2.0",
    "Terjman-Large-v2": "BounharAbdelaziz/Terjman-Large-v2.0",
    "Terjman-Ultra-v2": "BounharAbdelaziz/Terjman-Ultra-v2.0",
    "Terjman-Supreme-v2": "BounharAbdelaziz/Terjman-Supreme-v2.0"
}

# Load environment token
TOKEN = os.environ['TOKEN']

# Preload models and tokenizers
def preload_models():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    print(f"[INFO] Using device: {device}")

    # Load Nano and Large models
    nano_large_models = {}
    for model_name in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
        print(f"[INFO] Loading {model_name}...")
        translator = pipeline(
            "translation",
            model=MODEL_PATHS[model_name],
            token=TOKEN,
            device=device if device.startswith("cuda") else -1
        )
        nano_large_models[model_name] = translator

    # Load Ultra and Supreme models
    ultra_supreme_models = {}
    for model_name in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
        print(f"[INFO] Loading {model_name}...")
        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATHS[model_name], token=TOKEN).to(device)
        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATHS[model_name], token=TOKEN)
        translator = pipeline(
            "translation",
            model=model,
            tokenizer=tokenizer,
            device=device if device.startswith("cuda") else -1,
            src_lang="eng_Latn",
            tgt_lang="ary_Arab"
        )
        ultra_supreme_models[model_name] = translator

    return nano_large_models, ultra_supreme_models

# Preload all models
nano_large_models, ultra_supreme_models = preload_models()

# Translation function for Nano and Large models
@spaces.GPU
def translate_nano_large(text, model_name):
    translator = nano_large_models[model_name]
    translated = translator(
        text,
        max_length=512,
        num_beams=4,
        no_repeat_ngram_size=3,
        early_stopping=True,
        do_sample=False,
        pad_token_id=translator.tokenizer.pad_token_id,
        bos_token_id=translator.tokenizer.bos_token_id,
        eos_token_id=translator.tokenizer.eos_token_id,
    )
    return translated[0]["translation_text"]

# Translation function for Ultra and Supreme models
@spaces.GPU
def translate_ultra_supreme(text, model_name):
    translator = ultra_supreme_models[model_name]
    translation = translator(text)[0]['translation_text']
    return translation

# Main translation function
def translate_text(text, model_choice):
    if model_choice in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
        return translate_nano_large(text, model_choice)
    elif model_choice in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
        return translate_ultra_supreme(text, model_choice)
    else:
        return "Invalid model selection."

# Gradio app
def gradio_app():
    with gr.Blocks() as app:
        gr.Markdown("# 🇲🇦 Terjman-v2")
        gr.Markdown("Choose a model and enter the English text you want to translate to Moroccan Darija.")

        model_choice = gr.Dropdown(
            label="Select Model",
            choices=["Terjman-Nano-v2", "Terjman-Large-v2", "Terjman-Ultra-v2", "Terjman-Supreme-v2"],
            value="Terjman-Ultra-v2"
        )
        input_text = gr.Textbox(label="Input Text", placeholder="Enter text to translate...", lines=3)
        output_text = gr.Textbox(label="Translated Text", interactive=False, lines=3)
        translate_button = gr.Button("Translate")

        # Link input and output
        translate_button.click(
            fn=translate_text,
            inputs=[input_text, model_choice],
            outputs=output_text
        )

    return app

# Run the app
if __name__ == "__main__":
    app = gradio_app()
    app.launch()