kjarni-cli 0.1.9

Command-line interface for the Kjarni inference engine: embeddings, classification, reranking, generation and search
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
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

//! Summarization command using encoder-decoder models

use anyhow::{Result, anyhow};
use futures::{StreamExt, pin_mut};
use std::io::{self, Write};

use kjarni::summarizer::Summarizer;

use super::util::resolve_input;

pub async fn run(
    input: Option<&str>,
    model: &str,
    model_path: Option<&str>,
    min_length: Option<usize>,
    max_length: Option<usize>,
    num_beams: Option<usize>,
    length_penalty: Option<f32>,
    no_repeat_ngram: Option<usize>,
    greedy: bool,
    no_stream: bool,
    gpu: bool,
    quiet: bool,
) -> Result<()> {
    // Fail on an unknown model here, with suggestions, rather than letting the
    // builder report it without context.
    let _ = crate::commands::util::resolve_model(model, Some("seq2seq"))?;

    // Resolve input
    let text = resolve_input(input)?;
    if text.trim().is_empty() {
        return Err(anyhow!("Input text is empty. Nothing to summarize."));
    }

    // Build summarizer
    if model_path.is_some() {
        return Err(anyhow!("--model-path not yet implemented."));
    }

    let mut builder = Summarizer::builder(model);

    if gpu {
        builder = builder.gpu();
    } else {
        builder = builder.cpu();
    }

    if quiet {
        builder = builder.quiet();
    }

    if greedy {
        builder = builder.greedy();
    }

    if let Some(min) = min_length {
        builder = builder.min_length(min);
    }

    if let Some(max) = max_length {
        builder = builder.max_length(max);
    }

    if let Some(beams) = num_beams {
        builder = builder.num_beams(beams);
    }

    if let Some(penalty) = length_penalty {
        builder = builder.length_penalty(penalty);
    }

    if let Some(ngram) = no_repeat_ngram {
        builder = builder.no_repeat_ngram_size(ngram);
    }

    let summarizer = builder.build().await.map_err(|e| anyhow!("{}", e))?;

    // Show info
    if !quiet {
        eprintln!("Model: {}", summarizer.model_name());
        eprintln!("Device: {:?}", summarizer.device());
        eprintln!("Input: {} characters", text.len());
        eprintln!();
    }

    // Summarize
    if no_stream {
        let summary = summarizer
            .summarize(&text)
            .await
            .map_err(|e| anyhow!("Summarization failed: {}", e))?;
        println!("{}", summary.trim());
    } else {
        let stream = summarizer
            .stream(&text)
            .await
            .map_err(|e| anyhow!("Summarization failed: {}", e))?;

        pin_mut!(stream);

        let mut stdout = io::stdout();
        let mut generated_any = false;

        while let Some(token_result) = stream.next().await {
            let token = token_result.map_err(|e| anyhow!("Streaming error: {}", e))?;
            print!("{}", token);
            stdout.flush()?;
            generated_any = true;
        }

        if generated_any {
            println!();
        }
    }

    Ok(())
}