app.py

import gradio as gr
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
import os
import torch
import spaces
from datasets import Dataset
import time
import datetime

# 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 tokens
TOKEN = os.environ['TOKEN']

# Dataset configuration
DATASET_REPO = "BounharAbdelaziz/terjman-v2-live-translations"
# Number of translations to collect before pushing
BATCH_SIZE = 10  
# Time in seconds between pushes (1 hour)
UPDATE_INTERVAL = 3600

# Initialize dataset tracking
translations_buffer = []
last_push_time = time.time()

def preload_models():
    """ Preload models and tokenizers """
    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

def push_to_hf_dataset():
    """ Save translations in HF dataset for monitoring """
    global translations_buffer, last_push_time
    
    if not translations_buffer:
        return
    
    try:
        print(f"[INFO] Pushing {len(translations_buffer)} translations to Hugging Face dataset...")
        
        # Create dataset from buffer
        ds = Dataset.from_dict({
            "source_text": [item["source_text"] for item in translations_buffer],
            "translated_text": [item["translated_text"] for item in translations_buffer],
            "model_used": [item["model_used"] for item in translations_buffer],
            "timestamp": [item["timestamp"] for item in translations_buffer]
        })
        
        # Push to hub
        ds.push_to_hub(
            DATASET_REPO, 
            token=TOKEN,
            split=f"live_translations",
            private=True,
        )
        
        # Clear buffer and reset timer
        translations_buffer = []
        last_push_time = time.time()
        print("[INFO] Successfully pushed translations to Hugging Face dataset")
    
    except Exception as e:
        print(f"[ERROR] Failed to push dataset to Hugging Face: {str(e)}")

@spaces.GPU
def translate_nano_large(text, model_name):
    """ Translation function for Nano and Large models """
    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"]

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

def translate_text(text, model_choice):
    """ Main translation function """
    global translations_buffer, last_push_time
    
    # Skip empty text
    if not text or text.strip() == "":
        return "Please enter text to translate."
    
    # Perform translation
    if model_choice in ["Terjman-Nano-v2", "Terjman-Large-v2"]:
        translation = translate_nano_large(text, model_choice)
    elif model_choice in ["Terjman-Ultra-v2", "Terjman-Supreme-v2"]:
        translation = translate_ultra_supreme(text, model_choice)
    else:
        return "Invalid model selection."
    
    # Add to buffer
    translations_buffer.append({
        "source_text": text,
        "translated_text": translation,
        "model_used": model_choice,
        "timestamp": datetime.datetime.now().isoformat()
    })
    
    # Check if it's time to push to HF
    current_time = time.time()
    if len(translations_buffer) >= BATCH_SIZE or (current_time - last_push_time) >= UPDATE_INTERVAL:
        push_to_hf_dataset()
    
    return translation

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")
        
        # Status message -> hidden
        # status = gr.Markdown(f"Translations in buffer: 0")
        
        # Link input and output
        def translate_and_update_status(text, model):
            translation = translate_text(text, model)
            return translation #, f"Translations in buffer: {len(translations_buffer)} (Will push when reaching {BATCH_SIZE} or after {UPDATE_INTERVAL/3600} hours)"
        
        translate_button.click(
            fn=translate_and_update_status,
            inputs=[input_text, model_choice],
            outputs=[output_text] #, status] Status message -> hidden
        )
    
    return app

# Run the app
if __name__ == "__main__":
    # Register shutdown handler to save remaining translations
    import atexit
    atexit.register(push_to_hf_dataset)

    # Preload all models
    nano_large_models, ultra_supreme_models = preload_models()
    
    # Launch the app
    app = gradio_app()
    app.launch()