usertea commited on
Commit
0f04f89
·
1 Parent(s): abbbd7b

EchoScript : services\translation.py , Fix: services/translation.py now loads MarianMT directly via AutoModelForSeq2SeqLM + AutoTokenizer and calls model.generate() itself, bypassing the removed pipeline() task wrapper entirely. This is the path the official MarianMT docs use now too, so it's not a workaround you'll need to revisit on the next transformers release.

Browse files
Files changed (1) hide show
  1. services/translation.py +23 -7
services/translation.py CHANGED
@@ -34,24 +34,40 @@ class TranslationError(RuntimeError):
34
 
35
 
36
  @lru_cache(maxsize=None)
37
- def _load_pipeline(source_language: str, target_language: str):
38
- """Lazily load and cache a MarianMT pipeline for one language pair.
39
 
40
  Cached so repeated translations within a session don't reload a model
41
  from disk every time. Swapping the translation backend later (a
42
  different model, a hosted API, an offline engine like Argos) only
43
  requires changing this one function.
 
 
 
 
 
 
 
 
44
  """
45
- from transformers import pipeline # heavy import, deferred until needed
46
 
47
  model_name = f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"
48
  try:
49
- return pipeline("translation", model=model_name)
 
50
  except Exception as exc: # pragma: no cover - depends on model availability
51
  raise TranslationError(
52
  f"No translation model available for "
53
  f"'{source_language}' -> '{target_language}': {exc}"
54
  ) from exc
 
 
 
 
 
 
 
55
 
56
 
57
  class TranslationService:
@@ -67,16 +83,16 @@ class TranslationService:
67
  segments=list(transcript.segments),
68
  )
69
 
70
- translator = _load_pipeline(transcript.language, target_language)
71
 
72
  translated_segments = []
73
  for seg in transcript.segments:
74
  if not seg.text:
75
  translated_segments.append(seg)
76
  continue
77
- result_text = translator(seg.text)[0]["translation_text"]
78
  translated_segments.append(
79
- Segment(index=seg.index, start=seg.start, end=seg.end, text=result_text.strip())
80
  )
81
 
82
  return Translation(
 
34
 
35
 
36
  @lru_cache(maxsize=None)
37
+ def _load_engine(source_language: str, target_language: str):
38
+ """Lazily load and cache a MarianMT model+tokenizer for one language pair.
39
 
40
  Cached so repeated translations within a session don't reload a model
41
  from disk every time. Swapping the translation backend later (a
42
  different model, a hosted API, an offline engine like Argos) only
43
  requires changing this one function.
44
+
45
+ Note: this loads the model/tokenizer directly via AutoModelForSeq2SeqLM
46
+ rather than transformers.pipeline("translation", ...). As of
47
+ transformers v5, the generic "translation" pipeline task was removed
48
+ entirely (see huggingface/transformers#43825) -- pipeline("translation")
49
+ now raises KeyError. Loading the model directly and calling
50
+ `model.generate()` is unaffected by that change and is the path the
51
+ transformers docs now show for MarianMT.
52
  """
53
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer # heavy import, deferred
54
 
55
  model_name = f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"
56
  try:
57
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
58
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
59
  except Exception as exc: # pragma: no cover - depends on model availability
60
  raise TranslationError(
61
  f"No translation model available for "
62
  f"'{source_language}' -> '{target_language}': {exc}"
63
  ) from exc
64
+ return tokenizer, model
65
+
66
+
67
+ def _run_translation(tokenizer, model, text: str) -> str:
68
+ inputs = tokenizer(text, return_tensors="pt", truncation=True)
69
+ generated = model.generate(**inputs, max_new_tokens=512)
70
+ return tokenizer.decode(generated[0], skip_special_tokens=True).strip()
71
 
72
 
73
  class TranslationService:
 
83
  segments=list(transcript.segments),
84
  )
85
 
86
+ tokenizer, model = _load_engine(transcript.language, target_language)
87
 
88
  translated_segments = []
89
  for seg in transcript.segments:
90
  if not seg.text:
91
  translated_segments.append(seg)
92
  continue
93
+ result_text = _run_translation(tokenizer, model, seg.text)
94
  translated_segments.append(
95
+ Segment(index=seg.index, start=seg.start, end=seg.end, text=result_text)
96
  )
97
 
98
  return Translation(