byte_pair_encoding.rs

use std::collections::HashMap;
use std::env;
use std::fs;

fn main() {
	let args: Vec<String> = env::args().collect();
    if args.len() < 3 {
        eprintln!("Usage: {} <filename> <vocab_size>", args[0]);
        return;
    }

    let filename = &args[1];
    let target_vocab_size: u32 = args[2].parse().expect("Vocab size must be a number");
    let text = fs::read_to_string(filename).expect("Could not read file");

    let mut tokens: Vec<u32> = Vec::with_capacity(text.len());
    for byte in text.bytes() {
        tokens.push(byte as u32);
    }

	let mut current_vocab: u32 = 256;
	println!("Starting token count: {}", tokens.len().to_string());
	println!("Starting encoding in Rust...");
	println!("|----------------------|");

	loop {
		let mut lookup: HashMap<u64, u32> = HashMap::new();
		let mut max_freq: u32 = 0;
		let mut max_freq_token: u64 = 0;

		for token_index in 0..(tokens.len() - 1) {
			let key: u64 = ((tokens[token_index] as u64) << 32) | tokens[token_index+1] as u64;
			*lookup.entry(key).or_insert(0) += 1;
		}

		for (key, count) in &lookup {
			if *count > max_freq {
				max_freq = *count;
				max_freq_token = *key;
			}
		}

		if max_freq < 2 {break};
		let new_token_id = current_vocab;
		current_vocab += 1;

		let tokenA: u32 = (max_freq_token >> 32) as u32;
		let tokenB: u32 = (max_freq_token & 0xFFFFFFFF) as u32;
		let mut new_tokens: Vec<u32> = Vec::with_capacity(tokens.len());

		println!(
            "vocab: {:<5} / {} | merging: ({}, {}) -> {} | count: {}",
            current_vocab, target_vocab_size, tokenA, tokenB, new_token_id, max_freq
        );

		let mut i: usize = 0;
		while i < tokens.len() {
			if i < tokens.len() - 1 && tokens[i] == tokenA && tokens[i+1] == tokenB {
				new_tokens.push(new_token_id);
				i += 2;
			} else {
				new_tokens.push(tokens[i]);
				i += 1;
			}
		}
		tokens = new_tokens;
		if current_vocab > target_vocab_size {break};
	}

	println!("Encoding Complete.");
	println!("Final token count: {}", tokens.len());
}