chore: simplification, cleanup

This commit is contained in:
swaits committed 2024-09-25 21:10:26 -06:00
1 parent 608562c0c3
commit 485181549b
3 files changed
+25 -47

No files matched your search

Generated
+1 -1
View File
@@ -81,7 +81,7 @@ dependencies = [
[[package]] [[package]]
name = "bpetok" name = "bpetok"
version = "0.1.0" version = "0.1.1"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"bpe-tokenizer", "bpe-tokenizer",
+6 -2
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "bpetok" name = "bpetok"
version = "0.1.0" version = "0.1.1"
edition = "2021" edition = "2021"
authors = ["Stephen Waits <steve@waits.net>"] authors = ["Stephen Waits <steve@waits.net>"]
description = "A simple CLI for tokenizing text input using Byte Pair Encoding (BPE)." description = "A simple CLI for tokenizing text input using Byte Pair Encoding (BPE)."
@@ -10,5 +10,9 @@ categories = ["command-line-utilities", "text-processing"]
[dependencies] [dependencies]
anyhow = "1.0.89" 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"] } clap = { version = "4.5.18", features = ["derive"] }
+18 -44
View File
@@ -5,12 +5,11 @@ use std::{
use anyhow::Result; use anyhow::Result;
use bpe_tokenizer::BytePairEncoder; use bpe_tokenizer::BytePairEncoder;
use clap::{Parser, ArgAction}; use clap::{ArgAction, Parser};
/// A simple CLI for tokenizing text input using Byte Pair Encoding (BPE). /// A simple CLI for tokenizing text input using Byte Pair Encoding (BPE).
/// ///
/// This application accepts text from stdin and emits BPE tokens to stdout. /// This application accepts text from stdin and emits BPE tokens to stdout.
///
#[derive(Parser)] #[derive(Parser)]
#[command(version, about, long_about, author)] #[command(version, about, long_about, author)]
#[group(multiple = false)] #[group(multiple = false)]
@@ -36,57 +35,32 @@ fn main() -> Result<()> {
// Parse CLI arguments // Parse CLI arguments
let cli = Cli::parse(); let cli = Cli::parse();
// Load vocabulary from file // Load vocabulary
let vocab = if let Some(vocab_path) = cli.vocab.as_deref() { let bpe = if let Some(vocab_path) = cli.vocab.as_deref() {
match BytePairEncoder::new_from_file(vocab_path.to_str().unwrap()) { BytePairEncoder::new_from_file(
Ok(vocab) => vocab, vocab_path
Err(err) => { .to_str()
eprintln!("Error loading vocabulary file: {}", err); .ok_or(anyhow::anyhow!("Invalid vocabulary path"))?,
std::process::exit(1); )?
}
}
} else if cli.small.unwrap_or(false) { } else if cli.small.unwrap_or(false) {
match BytePairEncoder::new_default_small() { 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) { } else if cli.large.unwrap_or(false) {
match BytePairEncoder::new_default_large() { BytePairEncoder::new_default_large()?
Ok(bpe) => bpe,
Err(err) => {
eprintln!("Error initializing default-large vocabulary file: {}", err);
std::process::exit(1);
}
}
} else { } else {
match BytePairEncoder::new_default_medium() { 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 // Create a buffer reader for stdin
let reader = BufReader::new(io::stdin().lock()); let reader = BufReader::new(io::stdin().lock());
// Iterate over the lines of stdin, simultaneously tokenizing and emitting to stdout // Iterate over the lines of stdin while tokenizing and emitting to stdout
for line in reader.lines() { for line in reader.lines().map_while(Result::ok) {
let line = line?; bpe.tokenize_iter(&line).for_each(|t| {
println!("{}", t);
// Tokenize each line });
for token in vocab.tokenize_iter(&line) { io::stdout().flush()?;
// Emit token
println!("{}", token);
// Ensure output is flushed immediately
io::stdout().flush()?;
}
} }
// Success
Ok(()) Ok(())
} }