SAPI4 / app.py
NeoPy's picture
Update app.py
fedc99c verified
Raw History Blame
1.88 kB
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()