chore: more simplification, more cleanup

This commit is contained in:
swaits committed 2024-09-25 21:33:25 -06:00
1 parent 485181549b
commit 768a1ef1c3
3 files changed
+21 -27

No files matched your search

Generated
+1 -1
View File
@@ -81,7 +81,7 @@ dependencies = [
[[package]] [[package]]
name = "bpetok" name = "bpetok"
version = "0.1.1" version = "0.1.2"
dependencies = [ dependencies = [
"anyhow", "anyhow",
"bpe-tokenizer", "bpe-tokenizer",
+1 -1
View File
@@ -1,6 +1,6 @@
[package] [package]
name = "bpetok" name = "bpetok"
version = "0.1.1" version = "0.1.2"
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)."
+19 -25
View File
@@ -3,7 +3,7 @@ use std::{
path::PathBuf, path::PathBuf,
}; };
use anyhow::Result; use anyhow::{Context, Result};
use bpe_tokenizer::BytePairEncoder; use bpe_tokenizer::BytePairEncoder;
use clap::{ArgAction, Parser}; use clap::{ArgAction, Parser};
@@ -13,18 +13,18 @@ use clap::{ArgAction, Parser};
#[derive(Parser)] #[derive(Parser)]
#[command(version, about, long_about, author)] #[command(version, about, long_about, author)]
#[group(multiple = false)] #[group(multiple = false)]
struct Cli { struct VocabularyOption {
/// Use built-in small vocabulary (100k). /// Use built-in small vocabulary (100k).
#[arg(short, long, action = ArgAction::SetTrue)] #[arg(short, long, action = ArgAction::SetTrue)]
small: Option<bool>, small: bool,
/// Use built-in medium vocabulary (320k) [default] /// Use built-in medium vocabulary (320k) [default]
#[arg(short, long, action = ArgAction::SetTrue, default_value = "true")] #[arg(short, long, action = ArgAction::SetTrue, default_value_t = true)]
medium: Option<bool>, medium: bool,
/// Use built-in large vocabulary (1M). /// Use built-in large vocabulary (1M).
#[arg(short, long, action = ArgAction::SetTrue)] #[arg(short, long, action = ArgAction::SetTrue)]
large: Option<bool>, large: bool,
/// Path to custom BPE vocabulary file. /// Path to custom BPE vocabulary file.
#[arg(short, long, value_name = "FILE")] #[arg(short, long, value_name = "FILE")]
@@ -33,31 +33,25 @@ struct Cli {
fn main() -> Result<()> { fn main() -> Result<()> {
// Parse CLI arguments // Parse CLI arguments
let cli = Cli::parse(); let options = VocabularyOption::parse();
// Load vocabulary // Load vocabulary
let bpe = if let Some(vocab_path) = cli.vocab.as_deref() { let bpe = match options.vocab {
BytePairEncoder::new_from_file( Some(ref vocab_path) => BytePairEncoder::new_from_file(&vocab_path.to_string_lossy())?,
vocab_path _ if options.small => BytePairEncoder::new_default_small()?,
.to_str() _ if options.large => BytePairEncoder::new_default_large()?,
.ok_or(anyhow::anyhow!("Invalid vocabulary path"))?, _ => BytePairEncoder::new_default_medium()?, // Default case is medium
)?
} else if cli.small.unwrap_or(false) {
BytePairEncoder::new_default_small()?
} else if cli.large.unwrap_or(false) {
BytePairEncoder::new_default_large()?
} else {
BytePairEncoder::new_default_medium()?
}; };
// Create a buffer reader for stdin // Buffer reader for stdin
let reader = BufReader::new(io::stdin().lock()); let reader = BufReader::new(io::stdin().lock());
// Iterate over the lines of stdin while tokenizing and emitting to stdout // Process each line, tokenize, and print tokens
for line in reader.lines().map_while(Result::ok) { for line_result in reader.lines() {
bpe.tokenize_iter(&line).for_each(|t| { let line = line_result.context("Failed to read line from stdin")?;
println!("{}", t); for token in bpe.tokenize_iter(&line) {
}); println!("{}", token);
}
io::stdout().flush()?; io::stdout().flush()?;
} }