StenTokenRuntime-rs/src/main.rs

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(())
}