102 lines
3.2 KiB
Rust
102 lines
3.2 KiB
Rust
use std::collections::HashMap;
|
|
|
|
use anyhow::{Context as _, Result, ensure};
|
|
use candle_core::{IndexOp as _, Tensor};
|
|
use candle_onnx::onnx::ModelProto;
|
|
use galaxy_completer::ParsedTokensSnapshot;
|
|
use prost::Message as _;
|
|
use tokenizers::Tokenizer;
|
|
|
|
use super::{ClassificationResult, Model};
|
|
use crate::InputClassifierDecisionSource;
|
|
|
|
pub struct InferenceRunner {
|
|
model: ModelProto,
|
|
tokenizer: Tokenizer,
|
|
}
|
|
|
|
impl InferenceRunner {
|
|
pub fn new(model: Model) -> Result<Self> {
|
|
Ok(Self {
|
|
model: Self::load_model(model)?,
|
|
tokenizer: Self::load_tokenizer(model)?,
|
|
})
|
|
}
|
|
|
|
fn load_model(model: Model) -> Result<ModelProto> {
|
|
let model_bytes = model.bytes().ok_or_else(|| {
|
|
std::io::Error::new(std::io::ErrorKind::NotFound, "Model file not found")
|
|
})?;
|
|
let model = ModelProto::decode(model_bytes.as_ref())?;
|
|
Ok(model)
|
|
}
|
|
|
|
fn load_tokenizer(model: Model) -> Result<Tokenizer> {
|
|
let tokenizer_bytes = model.tokenizer_bytes().ok_or_else(|| {
|
|
std::io::Error::new(std::io::ErrorKind::NotFound, "Tokenizer file not found")
|
|
})?;
|
|
let tokenizer = Tokenizer::from_bytes(tokenizer_bytes).map_err(|e| anyhow::anyhow!(e))?;
|
|
Ok(tokenizer)
|
|
}
|
|
}
|
|
|
|
impl super::InferenceRunner for InferenceRunner {
|
|
fn run_inference(&self, input: &ParsedTokensSnapshot) -> Result<ClassificationResult> {
|
|
// Encode the input text into tokens.
|
|
let encoding = self
|
|
.tokenizer
|
|
.encode_fast(input.buffer_text.as_str(), true)
|
|
.map_err(|e| anyhow::anyhow!(e))?;
|
|
|
|
// For now, we'll do all inference on the CPU.
|
|
let device = candle_core::Device::Cpu;
|
|
|
|
let input_ids = Tensor::new(
|
|
encoding
|
|
.get_ids()
|
|
.iter()
|
|
.map(|&x| x as i64)
|
|
.collect::<Vec<_>>()
|
|
.as_slice(),
|
|
&device,
|
|
)
|
|
.context("failed to build input ids tensor")?;
|
|
let attention_mask = Tensor::new(
|
|
encoding
|
|
.get_attention_mask()
|
|
.iter()
|
|
.map(|&x| x as i64)
|
|
.collect::<Vec<_>>()
|
|
.as_slice(),
|
|
&device,
|
|
)
|
|
.context("failed to build attention mask tensor")?;
|
|
|
|
// Run inference.
|
|
let outputs = candle_onnx::simple_eval(
|
|
&self.model,
|
|
HashMap::from([
|
|
("input_ids".to_string(), input_ids.unsqueeze(0)?),
|
|
("attention_mask".to_string(), attention_mask.unsqueeze(0)?),
|
|
]),
|
|
)
|
|
.context("error evaluating the model")?;
|
|
|
|
let logits = outputs.get("logits").context("failed to get logits")?;
|
|
let probabilities = candle_nn::ops::softmax_last_dim(logits)
|
|
.context("failed to compute softmax")?
|
|
.i(0)
|
|
.context("failed to get first dimension")?
|
|
.to_vec1::<f32>()
|
|
.context("failed to convert softmax output to vec")?;
|
|
|
|
ensure!(probabilities.len() == 2, "expected 2 probabilities");
|
|
|
|
Ok(ClassificationResult {
|
|
p_ai: probabilities[0],
|
|
p_shell: probabilities[1],
|
|
source: InputClassifierDecisionSource::InputClassifier,
|
|
})
|
|
}
|
|
}
|