bossiara13/J.A.R.V.I.S-rust
1
1
Fork 0
Code Issues Pull requests Projects Releases 6 Packages Activity Actions

AI models shared registry + Code cleanup + Better async handling + Some fixes, etc

Browse source
This commit is contained in:
Priler 2026-02-18 21:08:48 +05:00
parent a8ff3442ff
commit 520b98143f
62 changed files with 1683 additions and 1239 deletions
Download patch file Download diff file Expand all files Collapse all files

244
crates/jarvis-core/src/slots/gliner.rs
View file

@ -2,123 +2,43 @@
// https://github.com/fbilhaut/gline-rs
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use once_cell::sync::OnceCell;
use parking_lot::Mutex;
use ndarray::Array;
use regex::Regex;
use tokenizers::Tokenizer;
use ort::value::Tensor;
pub mod structs;
use structs::GlinerModelInfo;
use std::fs;
use crate::commands::{SlotDefinition, SlotValue};
use crate::{APP_DIR, i18n};
use crate::models::gliner::GlinerModel;
// MODEL STATE
static MODEL: OnceCell<Arc<GlinerModel>> = OnceCell::new();
struct GlinerModel {
session: ort::session::Session,
tokenizer: Tokenizer,
splitter: Regex,
}
unsafe impl Send for GlinerModel {}
unsafe impl Sync for GlinerModel {}
static MODEL: OnceCell<Mutex<GlinerModel>> = OnceCell::new();
// GLiNER defaults (same as gline-rs Parameters::default())
// GLiNER defaults
const THRESHOLD: f32 = 0.3;
const MAX_WIDTH: usize = 12;
const MAX_LENGTH: usize = 512;
// applied after decoding
const MIN_CONFIDENCE: f32 = 0.4;
// word splitting regex (gline-rs RegexSplitter default)
const WORD_REGEX: &str = r"\w+(?:[-_]\w+)*|\S";
// INIT
pub fn init() -> Result<(), String> {
if MODEL.get().is_some() {
return Ok(());
}
let variant = crate::DB.get()
.map(|db| db.read().gliner_model.clone())
.unwrap_or_default();
let language = i18n::get_language();
let (model_dir, onnx_file) = if variant.is_empty() {
(select_model_dir(), "model.onnx".to_string())
} else {
crate::gliner_models::resolve_model(&variant, &language)
.unwrap_or_else(|| (select_model_dir(), "model.onnx".to_string()))
};
let model_path = model_dir.join("onnx").join(&onnx_file);
let tokenizer_path = model_dir.join("tokenizer.json");
info!("Loading GLiNER model from: {}, variant {}", model_dir.display(), variant);
let session = ort::session::Session::builder()
.map_err(|e| format!("Failed to create ort session builder: {}", e))?
.commit_from_file(&model_path)
.map_err(|e| format!("Failed to load ONNX model: {}", e))?;
let tokenizer = Tokenizer::from_file(&tokenizer_path)
.map_err(|e| format!("Failed to load tokenizer: {}", e))?;
let splitter = Regex::new(WORD_REGEX)
.map_err(|e| format!("Failed to compile word regex: {}", e))?;
MODEL.set(Mutex::new(GlinerModel { session, tokenizer, splitter }))
.map_err(|_| "GLiNER model already initialized".to_string())?;
info!("GLiNER model loaded");
pub fn init_with_model(model: Arc<GlinerModel>) -> Result<(), String> {
MODEL.set(model).map_err(|_| "GLiNER model already initialized".to_string())?;
info!("GLiNER slot extraction ready");
Ok(())
}
fn select_model_dir() -> PathBuf {
let base = APP_DIR.join("resources").join("models");
// word splitting
match i18n::get_language().as_str() {
"en" => {
let path = base.join("gliner_small-v2.1");
if path.exists() { return path; }
}
_ => {}
}
// multilingual (covers RU, UA, EN)
let multi = base.join("gliner_multi-v2.1");
if multi.exists() { return multi; }
// fallback
base.join("gliner_small-v2.1")
}
// WORD SPLITTING
struct WordToken {
struct WordToken<'a> {
start: usize,
end: usize,
text: String,
text: &'a str,
}
fn split_words(splitter: &Regex, text: &str, limit: Option<usize>) -> Vec<WordToken> {
fn split_words<'a>(text: &'a str, model: &GlinerModel, limit: Option<usize>) -> Vec<WordToken<'a>> {
let mut tokens = Vec::new();
for m in splitter.find_iter(text) {
for m in model.splitter.find_iter(text) {
tokens.push(WordToken {
start: m.start(),
end: m.end(),
text: m.as_str().to_string(),
text: m.as_str(),
});
if let Some(lim) = limit {
if tokens.len() >= lim { break; }
@ -127,7 +47,7 @@ fn split_words(splitter: &Regex, text: &str, limit: Option<usize>) -> Vec<WordTo
tokens
}
// PROMPT CONSTRUCTION
// prompt construction
//
// GLiNER prompt format:
// [<<ENT>>, label1_w1, label1_w2, <<ENT>>, label2_w1, ..., <<SEP>>, word1, word2, ..., wordN]
@ -137,20 +57,20 @@ fn build_prompt(entities: &[&str], words: &[WordToken]) -> (Vec<String>, usize)
for entity in entities {
prompt.push("<<ENT>>".to_string());
prompt.push(entity.to_string()); // whole string, no split
prompt.push(entity.to_string());
}
prompt.push("<<SEP>>".to_string());
let entities_len = prompt.len();
for w in words {
prompt.push(w.text.clone());
prompt.push(w.text.to_string());
}
(prompt, entities_len)
}
// ENCODING
// encoding
struct EncodedBatch {
input_ids: ndarray::Array2<i64>,
@ -161,8 +81,7 @@ struct EncodedBatch {
}
fn encode_single(
tokenizer: &Tokenizer,
_text: &str,
model: &GlinerModel,
entities: &[&str],
words: &[WordToken],
) -> Result<EncodedBatch, String> {
@ -174,7 +93,7 @@ fn encode_single(
let mut entity_tokens: usize = 0;
for (pos, word) in prompt.iter().enumerate() {
let encoding = tokenizer.encode(word.as_str(), false)
let encoding = model.tokenizer.encode(word.as_str(), false)
.map_err(|e| format!("Tokenizer encode error: {}", e))?;
let ids = encoding.get_ids().to_vec();
total_tokens += ids.len();
@ -184,13 +103,13 @@ fn encode_single(
word_encodings.push(ids);
}
// text_offset: index where text tokens start (after BOS + entity tokens)
let text_offset = entity_tokens + 1;
// DEBUG
debug!("GLiNER prompt ({} total, ent_len={}, text_offset={}):", prompt.len(), ent_len, text_offset);
for (i, (word, enc)) in prompt.iter().zip(word_encodings.iter()).enumerate() {
debug!(" [{}]{} '{}' -> {:?}", i, if i < ent_len { " ENT" } else { " TXT" }, word, enc);
if log::log_enabled!(log::Level::Debug) {
debug!("GLiNER prompt ({} total, ent_len={}, text_offset={}):", prompt.len(), ent_len, text_offset);
for (i, (word, enc)) in prompt.iter().zip(word_encodings.iter()).enumerate() {
debug!(" [{}]{} '{}' -> {:?}", i, if i < ent_len { " ENT" } else { " TXT" }, word, enc);
}
}
let mut input_ids = Array::zeros((1, total_tokens));
@ -205,18 +124,15 @@ fn encode_single(
attention_masks[[0, idx]] = 1;
idx += 1;
// encode each word - matching gline-rs idx-based logic exactly
for word_enc in word_encodings.iter() {
for (token_idx, &token_id) in word_enc.iter().enumerate() {
input_ids[[0, idx]] = token_id as i64;
attention_masks[[0, idx]] = 1;
// word mask: only for text tokens (past text_offset), first sub-token only
if idx >= text_offset && token_idx == 0 {
word_masks[[0, idx]] = word_id;
}
idx += 1;
}
// increment word_id for any word whose tokens end past text_offset
if idx >= text_offset {
word_id += 1;
}
@ -229,9 +145,11 @@ fn encode_single(
let mut text_lengths = Array::zeros((1, 1));
text_lengths[[0, 0]] = (text_word_count + 1) as i64;
debug!("GLiNER input_ids: {:?}", input_ids.as_slice().unwrap());
debug!("GLiNER word_masks: {:?}", word_masks.as_slice().unwrap());
debug!("GLiNER text_lengths: {}", text_word_count);
if log::log_enabled!(log::Level::Debug) {
debug!("GLiNER input_ids: {:?}", input_ids.as_slice().unwrap());
debug!("GLiNER word_masks: {:?}", word_masks.as_slice().unwrap());
debug!("GLiNER text_lengths: {}", text_word_count);
}
Ok(EncodedBatch {
input_ids,
@ -242,7 +160,7 @@ fn encode_single(
})
}
// SPAN TENSORS
// span tensors
fn make_span_tensors(num_words: usize, max_width: usize) -> (ndarray::Array3<i64>, ndarray::Array2<bool>) {
let num_spans = num_words * max_width;
@ -264,7 +182,7 @@ fn make_span_tensors(num_words: usize, max_width: usize) -> (ndarray::Array3<i64
(span_idx, span_mask)
}
// DECODE + GREEDY SEARCH
// decode + greedy search
fn sigmoid(x: f32) -> f32 {
1.0 / (1.0 + (-x).exp())
@ -323,56 +241,43 @@ fn decode_and_search(
}
spans.sort_unstable_by(|a, b| (a.start, a.end).cmp(&(b.start, b.end)));
greedy_flat(&spans)
greedy_flat(spans)
}
fn greedy_flat(spans: &[Entity]) -> Vec<Entity> {
if spans.is_empty() {
return Vec::new();
// takes ownership, filters in place - no cloning
fn greedy_flat(mut spans: Vec<Entity>) -> Vec<Entity> {
if spans.len() <= 1 {
return spans;
}
let mut result: Vec<Entity> = Vec::new();
let mut keep = vec![false; spans.len()];
let mut prev = 0usize;
let mut next = 1usize;
while next < spans.len() {
let p = &spans[prev];
let n = &spans[next];
for next in 1..spans.len() {
let no_overlap = spans[next].start >= spans[prev].end
|| spans[prev].start >= spans[next].end;
if n.start >= p.end || p.start >= n.end {
result.push(Entity {
text: p.text.clone(),
label: p.label.clone(),
prob: p.prob,
start: p.start,
end: p.end,
});
if no_overlap {
keep[prev] = true;
prev = next;
} else if p.prob < n.prob {
} else if spans[prev].prob < spans[next].prob {
prev = next;
}
next += 1;
}
keep[prev] = true;
let last = &spans[prev];
result.push(Entity {
text: last.text.clone(),
label: last.label.clone(),
prob: last.prob,
start: last.start,
end: last.end,
});
result
let mut idx = 0;
spans.retain(|_| { let k = keep[idx]; idx += 1; k });
spans
}
// PUBLIC API
// public extract API
pub fn extract(
text: &str,
slots: &HashMap<String, SlotDefinition>,
) -> Result<HashMap<String, SlotValue>, String> {
let mut model = MODEL.get().ok_or("GLiNER not initialized")?.lock();
let model = MODEL.get().ok_or("GLiNER not initialized")?;
let mut label_to_slots: HashMap<&str, Vec<&str>> = HashMap::new();
for (slot_name, def) in slots {
@ -392,12 +297,12 @@ pub fn extract(
debug!("GLiNER extract: text='{}', labels={:?}", text, labels);
let words = split_words(&model.splitter, text, Some(MAX_LENGTH));
let words = split_words(text, &model, Some(MAX_LENGTH));
if words.is_empty() {
return Ok(HashMap::new());
}
let encoded = encode_single(&model.tokenizer, text, &labels, &words)?;
let encoded = encode_single(&model, &labels, &words)?;
let (span_idx, span_mask) = make_span_tensors(encoded.num_words, MAX_WIDTH);
@ -408,7 +313,8 @@ pub fn extract(
let t_span_idx = Tensor::from_array(span_idx).map_err(|e| format!("tensor: {}", e))?;
let t_span_mask = Tensor::from_array(span_mask).map_err(|e| format!("tensor: {}", e))?;
let outputs = model.session.run(
let mut session = model.session.lock();
let outputs = session.run(
ort::inputs! {
"input_ids" => t_input_ids,
"attention_mask" => t_attn,
@ -425,27 +331,29 @@ pub fn extract(
let logits_shape: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
debug!("GLiNER logits shape: {:?}, data len: {}", logits_shape, logits_data.len());
let max_logit = logits_data.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
debug!("GLiNER max logit: {:.4}, sigmoid: {:.4}", max_logit, sigmoid(max_logit));
// debug dump - gated so sigmoid/loop don't run in release
if log::log_enabled!(log::Level::Debug) {
debug!("GLiNER logits shape: {:?}, data len: {}", logits_shape, logits_data.len());
let max_logit = logits_data.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
debug!("GLiNER max logit: {:.4}, sigmoid: {:.4}", max_logit, sigmoid(max_logit));
// dump all scores above 5%
let num_words = logits_shape.get(1).copied().unwrap_or(0);
let dim_mw = logits_shape.get(2).copied().unwrap_or(0);
let dim_e = logits_shape.get(3).copied().unwrap_or(0);
for start in 0..num_words {
for width in 0..dim_mw.min(num_words - start) {
for class_idx in 0..dim_e {
let flat_idx = start * dim_mw * dim_e + width * dim_e + class_idx;
if flat_idx < logits_data.len() {
let score = logits_data[flat_idx];
let prob = sigmoid(score);
if prob > 0.05 {
let end = start + width;
let w_start = if start < words.len() { &words[start].text } else { "?" };
let w_end = if end < words.len() { &words[end].text } else { "?" };
debug!(" span[{}..{}] '{}'->'{}' label={} score={:.3} prob={:.3}",
start, end, w_start, w_end, labels.get(class_idx).unwrap_or(&"?"), score, prob);
let num_words = logits_shape.get(1).copied().unwrap_or(0);
let dim_mw = logits_shape.get(2).copied().unwrap_or(0);
let dim_e = logits_shape.get(3).copied().unwrap_or(0);
for start in 0..num_words {
for width in 0..dim_mw.min(num_words - start) {
for class_idx in 0..dim_e {
let flat_idx = start * dim_mw * dim_e + width * dim_e + class_idx;
if flat_idx < logits_data.len() {
let score = logits_data[flat_idx];
let prob = sigmoid(score);
if prob > 0.05 {
let end = start + width;
let w_start = if start < words.len() { words[start].text } else { "?" };
let w_end = if end < words.len() { words[end].text } else { "?" };
debug!(" span[{}..{}] '{}'->'{}' label={} score={:.3} prob={:.3}",
start, end, w_start, w_end, labels.get(class_idx).unwrap_or(&"?"), score, prob);
}
}
}
}
@ -484,4 +392,4 @@ fn parse_slot_value(text: &str) -> SlotValue {
return SlotValue::Number(n);
}
SlotValue::Text(text.to_string())
}
}

7
crates/jarvis-core/src/slots/gliner/structs.rs
View file

@ -1,7 +0,0 @@
#[derive(Debug, Clone)]
pub struct GlinerModelInfo {
pub model_dir: String,
pub file_name: String,
pub display_name: String,
pub value: String,
}