Blame for wasm-src/src/lib.rs

use nedo_tokenizer::{
    SurfaceVocabulary, Tokenizer, TokenizerConfig, SURFACE_BYTE_BASE_ID, SURFACE_ENTRY_BASE_ID,
};
use serde_json::{json, Value};
use wasm_bindgen::prelude::*;

const SURFACE_VOCABULARY: &[u8] = include_bytes!("../assets/surface-vocab.bin");

#[derive(Clone, Copy)]
struct PieceSpan {
    id: u32,
    start: usize,
    end: usize,
    base_byte: bool,
}

#[wasm_bindgen]
pub fn tokenize_json(text: &str) -> String {
    match tokenize_inner(text) {
        Ok(value) => value,
        Err(error) => json!({"ok": false, "error": error}).to_string(),
    }
}

fn display_valid(value: &str) -> String {
    value
        .chars()
        .map(|character| match character {
            ' ' => "␠".to_owned(),
            '\t' => "⇥".to_owned(),
            '\r' => "␍".to_owned(),
            '\n' => "↵".to_owned(),
            other => other.to_string(),
        })
        .collect::<String>()
}

fn hex_bytes(bytes: &[u8]) -> String {
    bytes
        .iter()
        .map(|byte| format!("{byte:02X}"))
        .collect::<Vec<_>>()
        .join(" ")
}

fn visual_segments(text: &str, pieces: &[PieceSpan]) -> Vec<Value> {
    let mut chars = Vec::new();
    let mut piece_cursor = 0_usize;
    for (start, character) in text.char_indices() {
        let end = start + character.len_utf8();
        while piece_cursor < pieces.len() && pieces[piece_cursor].end <= start {
            piece_cursor += 1;
        }
        let mut owners = Vec::new();
        let mut probe = piece_cursor;
        while probe < pieces.len() && pieces[probe].start < end {
            if pieces[probe].end > start {
                owners.push(probe);
            }
            probe += 1;
        }
        chars.push((start, end, character, owners));
    }

    let mut output = Vec::new();
    let mut index = 0_usize;
    while index < chars.len() {
        let (start, mut end, character, owners) = &chars[index];
        if owners.len() == 1 {
            let owner = owners[0];
            let mut text_value = character.to_string();
            let mut next = index + 1;
            while next < chars.len() {
                let (_, next_end, next_character, next_owners) = &chars[next];
                if next_owners.as_slice() != [owner] {
                    break;
                }
                text_value.push(*next_character);
                end = *next_end;
                next += 1;
            }
            let piece = pieces[owner];
            output.push(json!({
                "start": *start,
                "end": end,
                "text": text_value,
                "display": display_valid(&text_value),
                "token_indexes": [owner],
                "token_ids": [piece.id],
                "split_utf8": false,
                "base_byte": piece.base_byte,
            }));
            index = next;
            continue;
        }

        let token_ids = owners.iter().map(|&owner| pieces[owner].id).collect::<Vec<_>>();
        let base_byte = owners.iter().any(|&owner| pieces[owner].base_byte);
        let text_value = character.to_string();
        output.push(json!({
            "start": *start,
            "end": end,
            "text": text_value,
            "display": display_valid(&text_value),
            "token_indexes": owners,
            "token_ids": token_ids,
            "split_utf8": true,
            "base_byte": base_byte,
        }));
        index += 1;
    }
    output
}

fn tokenize_inner(text: &str) -> Result<String, String> {
    let tokenizer = Tokenizer::embedded(TokenizerConfig::default()).map_err(|e| e.to_string())?;
    let vocabulary = SurfaceVocabulary::from_bytes(SURFACE_VOCABULARY).map_err(|e| e.to_string())?;
    let document = tokenizer
        .tokenize(text.as_bytes().to_vec())
        .map_err(|e| e.to_string())?;
    let encoded = vocabulary
        .encode_document(&document, false)
        .map_err(|e| e.to_string())?;
    let raw = document.raw();
    let decoded = vocabulary
        .decode_ids(&encoded.ids)
        .map_err(|e| e.to_string())?;
    let mut pieces = Vec::<Value>::new();
    let mut piece_spans = Vec::<PieceSpan>::new();
    let mut byte_cursor = 0_usize;
    let mut base_byte_tokens = 0_usize;
    let mut learned_tokens = 0_usize;

    for (&id, &length) in encoded.ids.iter().zip(&encoded.lengths) {
        if length == 0 {
            continue;
        }
        let start = byte_cursor;
        let end = start
            .checked_add(usize::from(length))
            .ok_or_else(|| "surface token byte offset overflow".to_owned())?;
        let bytes = raw
            .get(start..end)
            .ok_or_else(|| "surface token span outside input".to_owned())?;
        byte_cursor = end;
        let raw_id = u32::from(id);
        let base_byte = (SURFACE_BYTE_BASE_ID..SURFACE_ENTRY_BASE_ID).contains(&raw_id);
        let kind = if base_byte {
            base_byte_tokens += 1;
            "base-byte"
        } else {
            learned_tokens += 1;
            "learned"
        };
        let text_piece = std::str::from_utf8(bytes).ok();
        pieces.push(json!({
            "id": raw_id,
            "kind": kind,
            "length": length,
            "start": start,
            "end": end,
            "text": text_piece,
            "display": text_piece.map(display_valid),
            "hex": hex_bytes(bytes),
        }));
        piece_spans.push(PieceSpan {
            id: raw_id,
            start,
            end,
            base_byte,
        });
    }
    if byte_cursor != raw.len() {
        return Err("surface token byte accounting differs from input".to_owned());
    }

    let visual = visual_segments(text, &piece_spans);
    Ok(json!({
        "ok": true,
        "source_commit": "8bf1b7ea881d11dd3c230209572b9267f7a68962",
        "vocabulary_kind": format!("{:?}", vocabulary.kind()),
        "input_bytes": raw.len(),
        "roundtrip": decoded.as_slice() == raw,
        "vocabulary_size": vocabulary.len(),
        "content_tokens": pieces.len(),
        "learned_tokens": learned_tokens,
        "base_byte_tokens": base_byte_tokens,
        "pieces": pieces,
        "visual": visual,
    })
    .to_string())
}

#[cfg(test)]
mod tests {
    fn parsed(text: &str) -> serde_json::Value {
        let value = super::tokenize_inner(text).expect("tokenizer must run");
        serde_json::from_str(&value).expect("valid json")
    }

    #[test]
    fn production_asset_keeps_root_hard_and_suffix_chain_soft() {
        let parsed = parsed("Ananızı öpeyim");
        assert_eq!(parsed["ok"], true);
        assert_eq!(parsed["roundtrip"], true);
        assert_eq!(parsed["vocabulary_size"], 32_000);
        assert_eq!(parsed["vocabulary_kind"], "RootSuffixByteBpe");
        let texts = parsed["pieces"]
            .as_array()
            .expect("pieces")
            .iter()
            .map(|piece| piece["text"].as_str().expect("valid Turkish piece"))
            .collect::<Vec<_>>();
        assert_eq!(texts, ["Ana", "nızı", " öp", "eyim"]);
    }

    #[test]
    fn readable_unicode_view_never_renders_hex_escape_fragments() {
        let parsed = parsed("Российская Федерация — 中华人民共和国 😀");
        assert_eq!(parsed["roundtrip"], true);
        let visual = parsed["visual"].as_array().expect("visual array");
        let reconstructed = visual
            .iter()
            .map(|segment| segment["text"].as_str().expect("valid unicode visual"))
            .collect::<String>();
        assert_eq!(reconstructed, "Российская Федерация — 中华人民共和国 😀");
        assert!(visual.iter().all(|segment| !segment["display"]
            .as_str()
            .unwrap_or_default()
            .contains("\\x")));
    }

    #[test]
    fn multi_space_run_remains_readable_and_exact() {
        let parsed = parsed("        print(x)");
        assert_eq!(parsed["roundtrip"], true);
        let visual = parsed["visual"].as_array().expect("visual array");
        assert!(visual
            .iter()
            .any(|segment| segment["display"] == "␠␠␠␠␠␠␠␠"));
    }
}