wordchipper-cli-util 0.9.2

Utilities for the Wordchipper CLI
Documentation 
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

use std::sync::Arc;

use wordchipper::{
    Tokenizer,
    UnifiedTokenVocab,
    disk_cache::WordchipperDiskCache,
};

/// Model selector arg group.
#[derive(clap::Args, Debug)]
#[group(required = true, multiple = false)]
pub struct ModelSelectorArgs {
    /// Model to use for encoding.
    #[arg(long, default_value = "openai:r50k_base")]
    model: String,
}

impl ModelSelectorArgs {
    /// Get the model name.
    pub fn model(&self) -> &str {
        &self.model
    }

    /// Load the vocabulary.
    pub fn load_vocab(
        &self,
        disk_cache: &mut WordchipperDiskCache,
    ) -> Result<Arc<UnifiedTokenVocab<u32>>, Box<dyn std::error::Error>> {
        let vocab = wordchipper::load_vocab(self.model(), disk_cache)?
            .vocab()
            .clone();

        Ok(vocab)
    }

    /// Load the tokenizer.
    pub fn load_tokenizer(
        &self,
        disk_cache: &mut WordchipperDiskCache,
    ) -> Result<Arc<Tokenizer<u32>>, Box<dyn std::error::Error>> {
        let vocab = self.load_vocab(disk_cache)?;
        let tokenizer = wordchipper::TokenizerOptions::default().build(vocab);
        Ok(tokenizer)
    }
}