File size: 2,762 Bytes
367eee8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | use anyhow::{bail, Context, Result};
use std::env;
use std::io::Read;
use stentoken_runtime::StenToken;
fn usage() -> ! {
eprintln!("Usage:\n stentoken-runtime encode <bundle-dir> <text>\n stentoken-runtime decode <bundle-dir> <comma-separated-ids>\n stentoken-runtime encode-json <bundle-dir> < JSON-array-of-strings\n stentoken-runtime decode-json <bundle-dir> < JSON-array-of-ID-arrays");
std::process::exit(2);
}
fn main() -> Result<()> {
let mut args = env::args().skip(1);
let command = args.next().unwrap_or_else(|| usage());
let bundle = args.next().unwrap_or_else(|| usage());
let tokenizer =
StenToken::from_bundle(&bundle).with_context(|| format!("loading bundle {bundle:?}"))?;
match command.as_str() {
"encode" => {
let text = args.next().unwrap_or_else(|| usage());
if args.next().is_some() {
usage();
}
println!("{}", serde_json::to_string(&tokenizer.encode(&text)?)?);
}
"decode" => {
let encoded = args.next().unwrap_or_else(|| usage());
if args.next().is_some() {
usage();
}
let ids = if encoded.is_empty() {
Vec::new()
} else {
encoded
.split(',')
.map(|value| {
value
.trim()
.parse::<u32>()
.context("token IDs must be unsigned integers")
})
.collect::<Result<Vec<_>>>()?
};
println!("{}", tokenizer.decode(&ids));
}
"encode-json" => {
if args.next().is_some() {
usage();
}
let mut input = String::new();
std::io::stdin().read_to_string(&mut input)?;
let texts: Vec<String> =
serde_json::from_str(&input).context("stdin must be a JSON array of strings")?;
println!(
"{}",
serde_json::to_string(&tokenizer.encode_batch(&texts)?)?
);
}
"decode-json" => {
if args.next().is_some() {
usage();
}
let mut input = String::new();
std::io::stdin().read_to_string(&mut input)?;
let batches: Vec<Vec<u32>> =
serde_json::from_str(&input).context("stdin must be a JSON array of ID arrays")?;
let decoded: Vec<String> = batches.iter().map(|ids| tokenizer.decode(ids)).collect();
println!("{}", serde_json::to_string(&decoded)?);
}
_ => bail!("unknown command {command:?}"),
}
Ok(())
}
|