diff --git a/Cargo.lock b/Cargo.lock index 86d9e9b..6fcfde5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -81,7 +81,7 @@ dependencies = [ [[package]] name = "bpetok" -version = "0.1.0" +version = "0.1.1" dependencies = [ "anyhow", "bpe-tokenizer", diff --git a/Cargo.toml b/Cargo.toml index 18c3260..5657795 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "bpetok" -version = "0.1.0" +version = "0.1.1" edition = "2021" authors = ["Stephen Waits "] description = "A simple CLI for tokenizing text input using Byte Pair Encoding (BPE)." @@ -10,5 +10,9 @@ categories = ["command-line-utilities", "text-processing"] [dependencies] anyhow = "1.0.89" -bpe-tokenizer = { version = "0.1.4", features = ["default-small", "default-medium", "default-large"] } +bpe-tokenizer = { version = "0.1.4", features = [ + "default-small", + "default-medium", + "default-large", +] } clap = { version = "4.5.18", features = ["derive"] } diff --git a/src/main.rs b/src/main.rs index 17ec4fc..a1390bf 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,12 +5,11 @@ use std::{ use anyhow::Result; use bpe_tokenizer::BytePairEncoder; -use clap::{Parser, ArgAction}; +use clap::{ArgAction, Parser}; /// 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)] @@ -36,57 +35,32 @@ 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); - } - } + // Load vocabulary + let bpe = if let Some(vocab_path) = cli.vocab.as_deref() { + BytePairEncoder::new_from_file( + vocab_path + .to_str() + .ok_or(anyhow::anyhow!("Invalid vocabulary path"))?, + )? } 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); - } - } + BytePairEncoder::new_default_small()? } 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); - } - } + BytePairEncoder::new_default_large()? } else { - match BytePairEncoder::new_default_medium() { - Ok(bpe) => bpe, - Err(err) => { - eprintln!("Error initializing default-medium vocabulary file: {}", err); - std::process::exit(1); - } - } + BytePairEncoder::new_default_medium()? }; // 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()?; - } + // Iterate over the lines of stdin while tokenizing and emitting to stdout + for line in reader.lines().map_while(Result::ok) { + bpe.tokenize_iter(&line).for_each(|t| { + println!("{}", t); + }); + io::stdout().flush()?; } + // Success Ok(()) }