import gradio as gr
from gradio import Server
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Load model and tokenizer
MODEL_NAME = "alirezamsh/small100"
print("Loading model...")
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME)
print("Model loaded!")


app = Server() 

@app.api(name="translate")
def translate(text, prompt):
    """
    Translate text using T5 with a task prefix.
    Note: t5-small is trained mainly on English tasks and has limited
    multilingual translation ability. See note at the bottom.
    """
    if not text.strip():
        return "Please enter some text to translate."

    tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=512,
            num_beams=4,
            early_stopping=True,
        )

    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result




# ─────────────────────────────────────────────
# 5. LAUNCH
# ─────────────────────────────────────────────
app.launch()