Blame for src/tokens.rs

//! BPE byte-level próprio do KHTST_RUST (contrato do `dados/tokenizador.py`
//! do v23: vocabulário próprio, especiais reservados). Sem dependências
//! externas; treino incremental por contagem de pares com índice de posições.
//!
//! IDs especiais: 0=<pad> 1=<bos> 2=<eos> 3=<unk>; bytes 4..260; merges a partir de 260.

use std::collections::{HashMap, HashSet};

pub const PAD: u32 = 0;
pub const BOS: u32 = 1;
pub const EOS: u32 = 2;
pub const UNK: u32 = 3;
pub const BASE_BYTES: u32 = 4;
const ESPECIAIS: [&str; 4] = ["<pad>", "<bos>", "<eos>", "<unk>"];

#[derive(Debug, Clone, Default)]
pub struct Bpe {
    /// token por id (string; bytes crus para os 256 iniciais)
    pub vocab: Vec<String>,
    /// merges na ordem de aprendizado: (id_a, id_b) -> id_novo implícito = 260+índice
    pub merges: Vec<(u32, u32)>,
    rank: HashMap<(u32, u32), u32>,
    id_por_token: HashMap<String, u32>,
    cache: HashMap<String, Vec<u32>>,
}

impl Bpe {
    pub fn novo() -> Self {
        let mut b = Bpe::default();
        for (i, s) in ESPECIAIS.iter().enumerate() {
            b.id_por_token.insert(s.to_string(), i as u32);
            b.vocab.push(s.to_string());
        }
        for i in 0..256u32 {
            let s = String::from_utf8_lossy(&[i as u8]).to_string();
            b.id_por_token.insert(s.clone(), BASE_BYTES + i);
            b.vocab.push(s);
        }
        b
    }

    pub fn n_vocab(&self) -> usize {
        self.vocab.len()
    }

    fn id_merge(&self, indice_merge: usize) -> u32 {
        260 + indice_merge as u32
    }

    /// Treina BPE em `textos` até `vocab_alvo` tokens (inclui especiais + bytes).
    /// Contagem incremental: cada merge só visita as posições do par afetado.
    pub fn treinar(textos: &[String], vocab_alvo: usize) -> Self {
        let mut b = Bpe::novo();
        // palavras (por espaço em branco; o espaço separa palavras — sem marcador final)
        let mut contagem_palavra: HashMap<String, u32> = HashMap::new();
        for t in textos {
            for w in t.split_whitespace() {
                *contagem_palavra.entry(w.to_string()).or_insert(0) += 1;
            }
        }
        // estado: sequências por palavra + índice global de pares
        let mut seqs: Vec<Vec<u32>> = vec![];
        let mut pesos: Vec<u32> = vec![];
        for (w, c) in &contagem_palavra {
            seqs.push(w.as_bytes().iter().map(|&x| BASE_BYTES + x as u32).collect());
            pesos.push(*c);
        }
        // contagem de pares e posições
        let mut cont: HashMap<(u32, u32), u64> = HashMap::new();
        let mut pos: HashMap<(u32, u32), Vec<(usize, usize)>> = HashMap::new();
        for (si, s) in seqs.iter().enumerate() {
            for j in 0..s.len().saturating_sub(1) {
                let p = (s[j], s[j + 1]);
                *cont.entry(p).or_insert(0) += pesos[si] as u64;
                pos.entry(p).or_default().push((si, j));
            }
        }
        let n_merges_alvo = vocab_alvo.saturating_sub(b.vocab.len());
        for k in 0..n_merges_alvo {
            // melhor par por (frequência, ordem de criação implícita)
            let melhor = cont.iter().filter(|(_, &c)| c > 0).max_by_key(|(p, &c)| (c, std::cmp::Reverse(*p)));
            let (par, _) = match melhor {
                Some((p, c)) => (*p, *c),
                None => break,
            };
            let novo_id = b.id_merge(b.merges.len());
            let concatenado = {
                let a = b.vocab.get(par.0 as usize).cloned().unwrap_or_default();
                let c = b.vocab.get(par.1 as usize).cloned().unwrap_or_default();
                format!("{a}{c}")
            };
            // aplica o merge nas posições registradas (por baixo para índices estáveis)
            let mut ocupadas: HashSet<(usize, usize)> = HashSet::new();
            let mut afetadas: Vec<(u32, u32)> = vec![];
            if let Some(posicoes) = pos.get(&par).cloned() {
                for (si, j) in posicoes {
                    if ocupadas.contains(&(si, j)) {
                        continue;
                    }
                    let s = &mut seqs[si];
                    if j + 1 >= s.len() || (s[j], s[j + 1]) != par {
                        continue;
                    }
                    // remove pares vizinhos antigos
                    if j > 0 {
                        let esq = (s[j - 1], s[j]);
                        if let Some(c) = cont.get_mut(&esq) {
                            *c = c.saturating_sub(pesos[si] as u64);
                        }
                        afetadas.push(esq);
                    }
                    if j + 2 < s.len() {
                        let dir = (s[j + 1], s[j + 2]);
                        if let Some(c) = cont.get_mut(&dir) {
                            *c = c.saturating_sub(peso_como_u64(pesos[si]));
                        }
                        afetadas.push(dir);
                    }
                    s[j] = novo_id;
                    s.remove(j + 1);
                    ocupadas.insert((si, j));
                    ocupadas.insert((si, j.saturating_sub(1)));
                    // novos pares
                    if j > 0 {
                        let esq = (s[j - 1], novo_id);
                        *cont.entry(esq).or_insert(0) += pesos[si] as u64;
                        pos.entry(esq).or_default().push((si, j - 1));
                    }
                    if j + 1 < s.len() {
                        let dir = (novo_id, s[j + 1]);
                        *cont.entry(dir).or_insert(0) += pesos[si] as u64;
                        pos.entry(dir).or_default().push((si, j));
                    }
                }
            }
            let _ = afetadas;
            cont.remove(&par);
            pos.remove(&par);
            b.merges.push(par);
            b.vocab.push(concatenado.clone());
            b.id_por_token.insert(concatenado, novo_id);
            if k % 512 == 0 {
                log::info!("bpe: merge {k}/{} (vocab {})", n_merges_alvo, b.vocab.len());
            }
        }
        b
    }

    /// Codifica uma palavra (com cache).
    fn codificar_palavra(&mut self, w: &str) -> Vec<u32> {
        if let Some(v) = self.cache.get(w) {
            return v.clone();
        }
        let mut s: Vec<u32> = w.as_bytes().iter().map(|&x| BASE_BYTES + x as u32).collect();
        loop {
            // par de menor rank
            let melhor = s.windows(2).map(|p| (p[0], p[1])).filter_map(|p| self.rank.get(&p).map(|r| (r, p))).min_by_key(|(r, _)| *r);
            match melhor {
                Some((_, (a, b))) => {
                    let alvo = self.rank[&(a, b)] as usize;
                    let novo = self.id_merge(alvo);
                    let mut out = vec![];
                    let mut i = 0;
                    while i < s.len() {
                        if i + 1 < s.len() && s[i] == a && s[i + 1] == b {
                            out.push(novo);
                            i += 2;
                        } else {
                            out.push(s[i]);
                            i += 1;
                        }
                    }
                    s = out;
                }
                None => break,
            }
            if s.len() <= 1 {
                break;
            }
        }
        self.cache.insert(w.to_string(), s.clone());
        s
    }

    pub fn reconstruir_indices(&mut self) {
        self.rank.clear();
        for (i, (a, b)) in self.merges.iter().enumerate() {
            self.rank.insert((*a, *b), i as u32);
        }
    }

    pub fn codificar(&mut self, texto: &str) -> Vec<u32> {
        let espaco = BASE_BYTES + b' ' as u32;
        let mut ids = vec![BOS];
        for w in texto.split_whitespace() {
            ids.extend(self.codificar_palavra(w));
            ids.push(espaco);
        }
        ids.push(EOS);
        ids
    }

    pub fn decodificar(&self, ids: &[u32]) -> String {
        let mut out = String::new();
        for &i in ids {
            match i {
                PAD | BOS | EOS => {}
                _ => {
                    if let Some(s) = self.vocab.get(i as usize) {
                        if !ESPECIAIS.contains(&s.as_str()) {
                            out.push_str(s);
                        }
                    }
                }
            }
        }
        out.trim().to_string()
    }

    pub fn para_json(&self) -> serde_json::Value {
        serde_json::json!({
            "vocab": self.vocab,
            "merges": self.merges,
        })
    }

    pub fn de_json(v: &serde_json::Value) -> Option<Self> {
        let vocab: Vec<String> = v.get("vocab")?.as_array()?.iter().map(|x| x.as_str().unwrap_or_default().to_string()).collect();
        let merges: Vec<(u32, u32)> = v.get("merges")?.as_array()?.iter().filter_map(|x| {
            let a = x.get(0)?.as_u64()? as u32;
            let b = x.get(1)?.as_u64()? as u32;
            Some((a, b))
        }).collect();
        let mut b = Bpe { vocab, merges, ..Default::default() };
        for (i, s) in b.vocab.iter().enumerate() {
            b.id_por_token.insert(s.clone(), i as u32);
        }
        b.reconstruir_indices();
        Some(b)
    }
}

fn peso_como_u64(p: u32) -> u64 {
    p as u64
}

#[cfg(test)]
mod testes {
    use super::*;

    #[test]
    fn treina_e_recodifica() {
        let textos: Vec<String> = ["o gato come o peixe e o cão come o osso",
             "o gato mira o peixe e o cão mira o osso",
             "peixe osso gato cão o e come mira"].iter().map(|s| s.to_string()).collect();
        let mut b = Bpe::treinar(&textos, 300);
        let ids = b.codificar("o gato come o peixe");
        assert_eq!(ids.first(), Some(&BOS));
        assert_eq!(ids.last(), Some(&EOS));
        let texto = b.decodificar(&ids);
        assert_eq!(texto, "o gato come o peixe");
        assert!(b.n_vocab() >= 260);
    }
}