Skip to content

fix(nllb): refactor NLLB translator for correctness and performance#1158

Open
zachelnet wants to merge 1 commit into
zyddnys:mainfrom
zachelnet:pr/fix-nllb-pipeline
Open

fix(nllb): refactor NLLB translator for correctness and performance#1158
zachelnet wants to merge 1 commit into
zyddnys:mainfrom
zachelnet:pr/fix-nllb-pipeline

Conversation

@zachelnet

Copy link
Copy Markdown

Summary

Fixes several correctness and performance bugs in nllb.py and NLLBBigTranslator.

Issues fixed

Issue Before After
pipeline() recreated per sentence new pipeline object every call (very slow) tokenizer+model used directly
Model not on device from_pretrained() → CPU, device ignored .to(device) after loading
'cpu' gets :0 appended cpu:0 (invalid) cpu stays as-is
Wrong language code UKR 'Ukrainian' (string, not a FLORES-200 code) 'ukr_Cyrl'
Wrong language code POL 'POL' 'PLK' (consistent with rest of project)
*.safetensors excluded from download newer HF models never downloaded removed from ignore_patterns
_check_downloaded misses safetensors only checked pytorch_model.bin checks both formats
logger.warn (deprecated) deprecated call logger.warning

Impact

  • Significant speedup (model/tokenizer reused across sentences)
  • Fixes silent CPU-only execution on GPU setups including AMD ROCm
  • Fixes broken Ukrainian translation (wrong FLORES-200 code)
  • Fixes broken Polish translation (wrong internal language key)

Co-authored-by: GitHub Copilot <github-copilot[bot]@users.noreply.github.com>

- Use tokenizer+model directly instead of recreating pipeline() per sentence
- Move model to device via .to(device) (was missing, broke ROCm/CUDA usage)
- Fix device normalization: skip ':0' suffix for 'cpu'
- Fix language code bugs: POL -> PLK, UKR: 'Ukrainian' -> 'ukr_Cyrl'
- Fix _check_downloaded: also check model.safetensors (newer HF format)
- Remove *.safetensors from download ignore_patterns (prevented model download)
- Return '' on unrecoverable failure instead of per-sentence retry loop
- Fix logger.warn -> logger.warning (deprecated)
@zachelnet zachelnet force-pushed the pr/fix-nllb-pipeline branch from b56b1e5 to 13ba227 Compare July 3, 2026 20:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant