Spaces:
Running
Running
| 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"); | |
| struct PieceSpan { | |
| id: u32, | |
| start: usize, | |
| end: usize, | |
| base_byte: bool, | |
| } | |
| 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()) | |
| } | |
| 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") | |
| } | |
| 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"]); | |
| } | |
| 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"))); | |
| } | |
| 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"] == "␠␠␠␠␠␠␠␠")); | |
| } | |
| } | |