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


# ─────────────────────────────────────────────
# 4. gr.Server — REST API + OPTIONAL SWAGGER UI
# ─────────────────────────────────────────────
app = Server(
    title="free Translator API",
    summary="free Translator API",
    version="1.0.0",
)


# Optional: plain FastAPI GET health-check route
@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_ID}

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