app.py

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# --- 1. Configuration ---
MODEL_ID = "Thermostatic/neuraltranslate-27b-mt-nah-es-v1.2-4bit"

# --- 2. Load Model and Tokenizer (once at startup) ---
# Load the model and tokenizer
# device_map="auto" will automatically place the model on the available GPU
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

print("Model and tokenizer loaded successfully.")

# --- 3. The Translation Function with GPU decorator ---
@spaces.GPU() # Request a GPU for 60 seconds per run
def translate(text_to_translate: str) -> str:
    """
    Translates Nahuatl text to Spanish using the pre-loaded model.
    """
    # Create the chat message format required by the model's template
    messages = [
        {"role": "user", "content": text_to_translate},
    ]

    # Apply the chat template and tokenize the input
    # The tokenizer will add the special tokens like [INST] and [/INST]
    model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)

    # Generate the translation
    # We can increase max_new_tokens for longer translations
    generated_ids = model.generate(model_inputs, max_new_tokens=256, do_sample=True)

    # Decode the generated tokens, skipping the prompt
    # generated_ids[0] gets the first sequence. The slice removes the input prompt tokens.
    decoded_output = tokenizer.decode(generated_ids[0][model_inputs.shape[1]:], skip_special_tokens=True)

    return decoded_output

# --- 4. Create the Gradio Interface ---
demo = gr.Interface(
    fn=translate,
    inputs=gr.Textbox(
        label="Enter Nahuatl Text",
        lines=5,
        placeholder="Cualli tonalli. ¿Quen otimotlaniltih?"
    ),
    outputs=gr.Textbox(
        label="Translated Spanish Text",
        lines=5,
        interactive=False # The user should not edit the output
    ),
    title="Nahuatl to Spanish Neural Translator",
    description=(
        "A simple UI for translating text from Nahuatl to Spanish using the "
        "Thermostatic/neuraltranslate-27b-mt-nah-es-v1.2-4bit model. "
        "This Space is running on a GPU thanks to the @spaces.GPU decorator."
    ),
    examples=[
        ["Tlazocamatih"],
        ["Cualli tonalli"],
        ["¿Quen tinemi?"],
        ["Nimitztlazohtla"]
    ],
    allow_flagging="never"
)

# --- 5. Launch the App ---
if __name__ == "__main__":
    demo.launch()