93 lines
2.7 KiB
Rust
93 lines
2.7 KiB
Rust
use std::{
|
|
io::{self, BufRead, BufReader, Write},
|
|
path::PathBuf,
|
|
};
|
|
|
|
use anyhow::Result;
|
|
use bpe_tokenizer::BytePairEncoder;
|
|
use clap::{Parser, ArgAction};
|
|
|
|
/// A simple CLI for tokenizing text input using Byte Pair Encoding (BPE).
|
|
///
|
|
/// This application accepts text from stdin and emits BPE tokens to stdout.
|
|
///
|
|
#[derive(Parser)]
|
|
#[command(version, about, long_about, author)]
|
|
#[group(multiple = false)]
|
|
struct Cli {
|
|
/// Use built-in small vocabulary (100k).
|
|
#[arg(short, long, action = ArgAction::SetTrue)]
|
|
small: Option<bool>,
|
|
|
|
/// Use built-in medium vocabulary (320k) [default]
|
|
#[arg(short, long, action = ArgAction::SetTrue, default_value = "true")]
|
|
medium: Option<bool>,
|
|
|
|
/// Use built-in large vocabulary (1M).
|
|
#[arg(short, long, action = ArgAction::SetTrue)]
|
|
large: Option<bool>,
|
|
|
|
/// Path to custom BPE vocabulary file.
|
|
#[arg(short, long, value_name = "FILE")]
|
|
vocab: Option<PathBuf>,
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
// Parse CLI arguments
|
|
let cli = Cli::parse();
|
|
|
|
// Load vocabulary from file
|
|
let vocab = if let Some(vocab_path) = cli.vocab.as_deref() {
|
|
match BytePairEncoder::new_from_file(vocab_path.to_str().unwrap()) {
|
|
Ok(vocab) => vocab,
|
|
Err(err) => {
|
|
eprintln!("Error loading vocabulary file: {}", err);
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
} else if cli.small.unwrap_or(false) {
|
|
match BytePairEncoder::new_default_small() {
|
|
Ok(bpe) => bpe,
|
|
Err(err) => {
|
|
eprintln!("Error initializing default-small vocabulary file: {}", err);
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
} else if cli.large.unwrap_or(false) {
|
|
match BytePairEncoder::new_default_large() {
|
|
Ok(bpe) => bpe,
|
|
Err(err) => {
|
|
eprintln!("Error initializing default-large vocabulary file: {}", err);
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
} else {
|
|
match BytePairEncoder::new_default_medium() {
|
|
Ok(bpe) => bpe,
|
|
Err(err) => {
|
|
eprintln!("Error initializing default-medium vocabulary file: {}", err);
|
|
std::process::exit(1);
|
|
}
|
|
}
|
|
};
|
|
|
|
// Create a buffer reader for stdin
|
|
let reader = BufReader::new(io::stdin().lock());
|
|
|
|
// Iterate over the lines of stdin, simultaneously tokenizing and emitting to stdout
|
|
for line in reader.lines() {
|
|
let line = line?;
|
|
|
|
// Tokenize each line
|
|
for token in vocab.tokenize_iter(&line) {
|
|
// Emit token
|
|
println!("{}", token);
|
|
|
|
// Ensure output is flushed immediately
|
|
io::stdout().flush()?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|