spark_model/model/nllb/
lang.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Translation language configuration for the served NLLB model.
4//!
5//! NLLB is not a causal chat model: the source must be encoded as
6//! `[src_lang] + subwords + </s>` and generation is seeded with
7//! `forced_bos = tgt_lang`. The language *strings* (`eng_Latn`, `gvn_Latn`, …)
8//! are resolved to token ids by the server's tokenizer (the model has no
9//! tokenizer), so this struct carries only the already-resolved ids plus the
10//! architectural special tokens.
11
12/// Resolved per-deployment translation tokens. `src_lang_id`/`tgt_lang_id` come
13/// from `--src-lang`/`--tgt-lang` (or a recipe default) resolved through the
14/// tokenizer at serve start; the rest are M2M-100/NLLB architectural constants.
15#[derive(Debug, Clone, Copy)]
16pub struct NllbLang {
17    /// Source-language prefix token prepended to the encoder input.
18    pub src_lang_id: u32,
19    /// Target-language token forced as the first decoded token (`forced_bos`).
20    pub tgt_lang_id: u32,
21    /// Decoder start token (M2M-100 convention: the eos id).
22    pub decoder_start_id: u32,
23    /// End-of-sequence / stop token.
24    pub eos_id: u32,
25    /// Padding token id (M2M-100 `padding_idx = 1`).
26    pub pad_id: u32,
27}
28
29impl NllbLang {
30    /// Format the raw source subword ids into the encoder input
31    /// `[src_lang] + tokens + </s>` using the deployment-default source language.
32    pub(super) fn encoder_input(&self, tokens: &[u32]) -> Vec<u32> {
33        self.encoder_input_with(self.src_lang_id, tokens)
34    }
35
36    /// Encoder input with an explicit (per-request) source-language token.
37    pub(super) fn encoder_input_with(&self, src_lang_id: u32, tokens: &[u32]) -> Vec<u32> {
38        let mut ids = Vec::with_capacity(tokens.len() + 2);
39        ids.push(src_lang_id);
40        ids.extend_from_slice(tokens);
41        ids.push(self.eos_id);
42        ids
43    }
44}