File size: 1,335 Bytes
e261077 12ccecd e261077 c690045 768990a fedc99c e261077 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | 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() |