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}