Spaces:
Running on Zero
Running on Zero
File size: 2,066 Bytes
88a619b 36fc86c 88a619b 9e83c68 88a619b 9e83c68 88a619b 9e83c68 88a619b 9e83c68 88a619b 9e83c68 88a619b | 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 45 46 47 48 49 50 51 52 53 54 55 | """CPU translation backend using billingsmoore/mlotsawa-ground-base
(a Tibetan->English seq2seq T5 model). Used whenever no OpenRouter API key is
available. The model is loaded once, lazily, on first use.
Generation uses the model's own task_specific_params["translation_bo_to_en"]
(prefix "translate Tibetan to English: ", num_beams=4, max_length=300) rather
than transformers' old pipeline("translation", ...) task, which was removed
in transformers>=5 (raises "Unknown task translation"). Calling generate()
directly with the model's documented settings works on any transformers
version that ships AutoModelForSeq2SeqLM/AutoTokenizer.
This backend ignores the editable translation prompt — it's a plain
seq2seq model, not an instruction-following LLM.
translate_batch is decorated with @spaces.GPU so this runs on HF ZeroGPU
Spaces (which refuse to boot a gradio-sdk app with no @spaces.GPU function
at all); the decorator is a no-op outside a ZeroGPU Space.
"""
import spaces
MODEL_ID = "billingsmoore/mlotsawa-ground-base"
_PREFIX = "translate Tibetan to English: "
_model = None
_tokenizer = None
def _load():
global _model, _tokenizer
if _model is None:
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
print(f"[INFO] Loading local CPU translation model: {MODEL_ID}")
_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
_model.eval()
return _model, _tokenizer
@spaces.GPU
def translate_batch(texts: list[str]) -> list[str]:
if not texts:
return []
import torch
model, tokenizer = _load()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
inputs = tokenizer(
[_PREFIX + t for t in texts],
return_tensors="pt", padding=True, truncation=True,
).to(device)
with torch.no_grad():
outputs = model.generate(**inputs, max_length=300, num_beams=4, early_stopping=True)
return [tokenizer.decode(o, skip_special_tokens=True) for o in outputs]
|