metrale_model_engine/model/nllb/
lang.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Translation language tokens for the served NLLB model.
4//!
5//! The encoder input is `[src_lang] + tokens + </s>`, and decoding is seeded with the
6//! target-language token (`forced_bos`). The server's tokenizer resolves language names
7//! to ids, so this struct carries only ids.
8//!
9//! Owner: model-engine.
10//! Invariants: none beyond the types.
11
12/// 2026-09-25: Resolved translation tokens for the deployment. `src_lang_id` and
13/// `tgt_lang_id` are the defaults from `--src-lang`/`--tgt-lang`; a request's own
14/// non-zero language ids override them. The factory sets `decoder_start_id` and
15/// `eos_id` to the config's `eos_token_id` and `pad_id` to 1.
16#[derive(Debug, Clone, Copy)]
17pub struct NllbLang {
18    /// 2026-09-25: Source-language token prepended to the encoder input.
19    pub src_lang_id: u32,
20    /// 2026-09-25: Target-language token forced as the first decoded token (`forced_bos`).
21    pub tgt_lang_id: u32,
22    /// 2026-09-25: Decoder start token, the first decoder input.
23    pub decoder_start_id: u32,
24    /// 2026-09-25: End-of-sequence token: appended to the encoder input, and it
25    /// ends a beam hypothesis.
26    pub eos_id: u32,
27    /// 2026-09-25: Padding token id; encoder positions skip it.
28    pub pad_id: u32,
29}
30
31impl NllbLang {
32    /// 2026-09-25: Format the raw source subword ids into the encoder input
33    /// `[src_lang] + tokens + </s>` using the deployment-default source language.
34    pub(super) fn encoder_input(&self, tokens: &[u32]) -> Vec<u32> {
35        self.encoder_input_with(self.src_lang_id, tokens)
36    }
37
38    /// 2026-09-25: Encoder input with an explicit (per-request) source-language token.
39    pub(super) fn encoder_input_with(&self, src_lang_id: u32, tokens: &[u32]) -> Vec<u32> {
40        let mut ids = Vec::with_capacity(tokens.len() + 2);
41        ids.push(src_lang_id);
42        ids.extend_from_slice(tokens);
43        ids.push(self.eos_id);
44        ids
45    }
46}