chore: more simplification, more cleanup
This commit is contained in:
3 files changed
+21
-27
No files matched your search
Generated
+1
-1
@@ -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
@@ -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
@@ -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()?;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user