DevEmmy commited on
Commit
aaf92dc
·
1 Parent(s): 80a7b2e

Drive the MT model directly: transformers 5 dropped the translation task

Browse files

pipeline('translation') raises KeyError on transformers 5, and gradio 6's
huggingface-hub 1.x requirement rules out going back to 4.x. Use
AutoModelForSeq2SeqLM with beam search instead - deterministic, which is the
point of moving off Gemini.

Files changed (1) hide show
  1. app.py +20 -3
app.py CHANGED
@@ -103,7 +103,14 @@ def _load(kind: str):
103
  device=DEVICE,
104
  )
105
  else:
106
- obj = pipeline("translation", model=MT_MODEL, token=HF_TOKEN, device=DEVICE)
 
 
 
 
 
 
 
107
  except Exception as e:
108
  _errors[kind] = f"{type(e).__name__}: {e}"
109
  raise
@@ -268,7 +275,17 @@ class TranslateIn(BaseModel):
268
 
269
 
270
  def _mt_impl(texts):
271
- return _load("mt")(texts, batch_size=MT_BATCH, truncation=True)
 
 
 
 
 
 
 
 
 
 
272
 
273
 
274
  _mt_gpu = _GPU(_mt_impl)
@@ -298,7 +315,7 @@ def translate(body: TranslateIn, x_service_token: str = Header(default=None)):
298
  if idx:
299
  res = _mt_run([texts[i] for i in idx])
300
  for i, r in zip(idx, res):
301
- out[i] = (r.get("translation_text") or "").strip()
302
  return {"translations": out}
303
 
304
 
 
103
  device=DEVICE,
104
  )
105
  else:
106
+ # transformers 5 dropped the "translation" pipeline task, so the
107
+ # seq2seq model is driven directly. Beam search, no sampling —
108
+ # the same cue gives the same English on every run, which is the
109
+ # whole reason for moving off Gemini.
110
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
111
+ tok = AutoTokenizer.from_pretrained(MT_MODEL, token=HF_TOKEN)
112
+ mdl = AutoModelForSeq2SeqLM.from_pretrained(MT_MODEL, token=HF_TOKEN)
113
+ obj = (tok, mdl.to(DEVICE).eval())
114
  except Exception as e:
115
  _errors[kind] = f"{type(e).__name__}: {e}"
116
  raise
 
275
 
276
 
277
  def _mt_impl(texts):
278
+ import torch
279
+ tok, mdl = _load("mt")
280
+ out = []
281
+ for start in range(0, len(texts), MT_BATCH):
282
+ chunk = texts[start:start + MT_BATCH]
283
+ enc = tok(chunk, return_tensors="pt", padding=True, truncation=True,
284
+ max_length=512).to(DEVICE)
285
+ with torch.inference_mode():
286
+ gen = mdl.generate(**enc, num_beams=4, max_new_tokens=256)
287
+ out.extend(tok.batch_decode(gen, skip_special_tokens=True))
288
+ return out
289
 
290
 
291
  _mt_gpu = _GPU(_mt_impl)
 
315
  if idx:
316
  res = _mt_run([texts[i] for i in idx])
317
  for i, r in zip(idx, res):
318
+ out[i] = (r or "").strip()
319
  return {"translations": out}
320
 
321