pub const BYTE_OFFSET: i64 = 3;
pub const PAD_ID: i64 = 0;
pub const EOS_ID: i64 = 1;
pub const UNK_ID: i64 = 2;
pub fn encode(text: &str) -> Vec<i64> {
let mut ids: Vec<i64> = text
.as_bytes()
.iter()
.map(|&b| i64::from(b) + BYTE_OFFSET)
.collect();
ids.push(EOS_ID);
ids
}
pub fn decode(token_ids: &[i64]) -> String {
let mut bytes = Vec::with_capacity(token_ids.len());
for &token in token_ids {
if token == EOS_ID {
break;
}
if token == PAD_ID || token == UNK_ID {
continue;
}
bytes.push((token - BYTE_OFFSET).rem_euclid(256) as u8);
}
String::from_utf8_lossy(&bytes).into_owned()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn encode_appends_offset_and_eos() {
assert_eq!(encode("a"), vec![100, EOS_ID]);
assert_eq!(encode(""), vec![EOS_ID]);
}
#[test]
fn encode_multibyte_is_per_byte() {
assert_eq!(encode("ខ"), vec![0xE1 + 3, 0x9E + 3, 0x81 + 3, EOS_ID]);
}
#[test]
fn decode_roundtrip_ascii_and_multibyte() {
let text = "ភាសា en ک";
assert_eq!(decode(&encode(text)), text);
}
#[test]
fn decode_stops_at_eos_and_skips_pad_unk() {
assert_eq!(decode(&[PAD_ID, 100, UNK_ID, 101, EOS_ID, 102]), "ab");
}
#[test]
fn decode_ids_below_offset_mod_256() {
assert_eq!(decode(&[324, EOS_ID]), "A");
}
}