92 lines
2.9 KiB
Rust
92 lines
2.9 KiB
Rust
use anyhow::{Result, ensure};
|
|
use itertools::Itertools as _;
|
|
use ort::{
|
|
execution_providers::CPUExecutionProvider, session::Session, tensor::ArrayExtensions as _,
|
|
value::Value,
|
|
};
|
|
use parking_lot::Mutex;
|
|
use tokenizers::Tokenizer;
|
|
use galaxy_completer::ParsedTokensSnapshot;
|
|
|
|
use super::ClassificationResult;
|
|
|
|
use super::Model;
|
|
|
|
pub struct InferenceRunner {
|
|
session: Mutex<Session>,
|
|
tokenizer: Tokenizer,
|
|
}
|
|
|
|
impl InferenceRunner {
|
|
pub fn new(model: Model) -> Result<Self> {
|
|
Ok(Self {
|
|
session: Self::init_session(model)?.into(),
|
|
tokenizer: Self::load_tokenizer(model)?,
|
|
})
|
|
}
|
|
|
|
fn init_session(model: Model) -> Result<Session> {
|
|
let model_bytes = model.bytes().ok_or_else(|| {
|
|
std::io::Error::new(std::io::ErrorKind::NotFound, "Model file not found")
|
|
})?;
|
|
let session = Session::builder()?
|
|
// For now, we'll do all inference on the CPU.
|
|
.with_execution_providers([CPUExecutionProvider::default().build()])?
|
|
.commit_from_memory(model_bytes.as_ref())?;
|
|
Ok(session)
|
|
}
|
|
|
|
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))?;
|
|
|
|
let input_ids = encoding.get_ids();
|
|
let attention_mask = encoding.get_attention_mask();
|
|
|
|
let input_ids = Value::from_array((
|
|
[1, input_ids.len()],
|
|
input_ids.iter().map(|&x| x as i64).collect_vec(),
|
|
))?;
|
|
let attention_mask = Value::from_array((
|
|
[1, attention_mask.len()],
|
|
attention_mask.iter().map(|&x| x as i64).collect_vec(),
|
|
))?;
|
|
|
|
let mut session = self.session.lock();
|
|
let outputs = session.run(ort::inputs![
|
|
"input_ids" => input_ids,
|
|
"attention_mask" => attention_mask,
|
|
])?;
|
|
|
|
let logits = &outputs[0];
|
|
|
|
let logits = logits.try_extract_array::<f32>()?;
|
|
let probabilities = logits.softmax(ndarray::Axis(1));
|
|
|
|
let probabilities = probabilities.view();
|
|
let probabilities = probabilities
|
|
.as_slice()
|
|
.ok_or_else(|| anyhow::anyhow!("failed to get probabilities"))?;
|
|
|
|
ensure!(probabilities.len() == 2, "expected 2 probabilities");
|
|
|
|
Ok(ClassificationResult {
|
|
p_ai: probabilities[0],
|
|
p_shell: probabilities[1],
|
|
})
|
|
}
|
|
}
|