| import gradio as gr |
| from gradio import Server |
| from transformers import T5ForConditionalGeneration, T5Tokenizer |
|
|
| |
| 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 |
|
|
|
|
|
|
|
|
| |
| |
| |
| app.launch() |
