AI models shared registry + Code cleanup + Better async handling + Some fixes, etc
This commit is contained in:
Priler
parent
a8ff3442ff
commit
520b98143f
62 changed files with 1683 additions and 1239 deletions
|
|
@ -1,79 +1,42 @@
|
|||
use parking_lot::Mutex;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::fs;
|
||||
|
||||
// use fastembed::{TextEmbedding, InitOptions, EmbeddingModel};
|
||||
use fastembed::{TextEmbedding, UserDefinedEmbeddingModel, TokenizerFiles, InitOptionsUserDefined, Pooling, QuantizationMode, OutputKey};
|
||||
use once_cell::sync::OnceCell;
|
||||
|
||||
use crate::commands::JCommandsList;
|
||||
use crate::i18n::get_language;
|
||||
use crate::{APP_CONFIG_DIR, APP_DIR, i18n};
|
||||
use crate::i18n;
|
||||
use crate::APP_CONFIG_DIR;
|
||||
use crate::models::embedding::EmbeddingModel;
|
||||
|
||||
static CLASSIFIER: OnceCell<Mutex<EmbeddingClassifier>> = OnceCell::new();
|
||||
// no outer Mutex needed - state is immutable after init.
|
||||
// the embedding model has its own internal Mutex.
|
||||
static CLASSIFIER: OnceCell<EmbeddingClassifierState> = OnceCell::new();
|
||||
|
||||
struct IntentVector {
|
||||
id: String,
|
||||
vector: Vec<f32>,
|
||||
}
|
||||
|
||||
struct EmbeddingClassifier {
|
||||
model: TextEmbedding,
|
||||
struct EmbeddingClassifierState {
|
||||
model: Arc<EmbeddingModel>,
|
||||
intents: Vec<IntentVector>,
|
||||
}
|
||||
|
||||
// model is Arc (Send+Sync), intents are read-only after init
|
||||
unsafe impl Send for EmbeddingClassifierState {}
|
||||
unsafe impl Sync for EmbeddingClassifierState {}
|
||||
|
||||
const CACHE_FILE: &str = "embedding_intents.json";
|
||||
const HASH_FILE: &str = "embedding_hash.txt";
|
||||
|
||||
pub fn init(commands: &[JCommandsList]) -> Result<(), String> {
|
||||
// init with a model loaded through the registry
|
||||
pub fn init_with_model(model: Arc<EmbeddingModel>, commands: &[JCommandsList]) -> Result<(), String> {
|
||||
if CLASSIFIER.get().is_some() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
info!("Initializing embedding model...");
|
||||
|
||||
// let mut model = TextEmbedding::try_new(
|
||||
// InitOptions::new(EmbeddingModel::AllMiniLML6V2).with_show_download_progress(true),
|
||||
// ).map_err(|e| format!("Failed to load embedding model: {}", e))?;
|
||||
|
||||
let model_dir;
|
||||
match i18n::get_language().as_str() {
|
||||
"en" => {
|
||||
// smaller model for English
|
||||
info!("Loading all-MiniLM-L6-v2 ...");
|
||||
model_dir = APP_DIR.join("resources").join("models").join("all-MiniLM-L6-v2");
|
||||
},
|
||||
_ => {
|
||||
// bigger model for any other languages (multilingual)
|
||||
info!("Loading paraphrase-multilingual-MiniLM-L12-v2-onnx-Q ...");
|
||||
model_dir = APP_DIR.join("resources").join("models").join("paraphrase-multilingual-MiniLM-L12-v2-onnx-Q");
|
||||
}
|
||||
}
|
||||
|
||||
// info!("{}", model_dir.display());
|
||||
|
||||
let user_model = UserDefinedEmbeddingModel {
|
||||
onnx_file: std::fs::read(model_dir.join("model.onnx"))
|
||||
.map_err(|e| format!("Failed to read model.onnx: {}", e))?,
|
||||
tokenizer_files: TokenizerFiles {
|
||||
tokenizer_file: std::fs::read(model_dir.join("tokenizer.json"))
|
||||
.map_err(|e| format!("Failed to read tokenizer.json: {}", e))?,
|
||||
config_file: std::fs::read(model_dir.join("config.json"))
|
||||
.map_err(|e| format!("Failed to read config.json: {}", e))?,
|
||||
special_tokens_map_file: std::fs::read(model_dir.join("special_tokens_map.json"))
|
||||
.map_err(|e| format!("Failed to read special_tokens_map.json: {}", e))?,
|
||||
tokenizer_config_file: std::fs::read(model_dir.join("tokenizer_config.json"))
|
||||
.map_err(|e| format!("Failed to read tokenizer_config.json: {}", e))?,
|
||||
},
|
||||
pooling: Some(Pooling::Mean),
|
||||
quantization: QuantizationMode::None,
|
||||
output_key: Some(OutputKey::ByName("last_hidden_state")),
|
||||
};
|
||||
|
||||
let mut model = TextEmbedding::try_new_from_user_defined(user_model, Default::default())
|
||||
.map_err(|e| format!("Failed to load embedding model: {}", e))?;
|
||||
|
||||
info!("Embedding model loaded");
|
||||
info!("Initializing embedding classifier...");
|
||||
|
||||
let current_hash = crate::commands::commands_hash(commands);
|
||||
let config_dir = APP_CONFIG_DIR.get().ok_or("Config dir not set")?;
|
||||
|
|
@ -90,7 +53,7 @@ pub fn init(commands: &[JCommandsList]) -> Result<(), String> {
|
|||
|
||||
let intents = if should_retrain {
|
||||
info!("Building intent vectors from commands...");
|
||||
let intents = build_intent_vectors(&mut model, commands)?;
|
||||
let intents = build_intent_vectors(&model, commands)?;
|
||||
|
||||
// cache to disk
|
||||
if let Ok(json) = serde_json::to_string(&intents_to_cache(&intents)) {
|
||||
|
|
@ -107,14 +70,14 @@ pub fn init(commands: &[JCommandsList]) -> Result<(), String> {
|
|||
|
||||
info!("Embedding classifier ready with {} intents", intents.len());
|
||||
|
||||
CLASSIFIER.set(Mutex::new(EmbeddingClassifier { model, intents }))
|
||||
.map_err(|_| "Classifier already set")?;
|
||||
CLASSIFIER.set(EmbeddingClassifierState { model, intents })
|
||||
.map_err(|_| "Classifier already set".to_string())?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn build_intent_vectors(
|
||||
model: &mut TextEmbedding,
|
||||
model: &EmbeddingModel,
|
||||
commands: &[JCommandsList],
|
||||
) -> Result<Vec<IntentVector>, String> {
|
||||
let lang = i18n::get_language();
|
||||
|
|
@ -129,7 +92,7 @@ fn build_intent_vectors(
|
|||
|
||||
let texts: Vec<&str> = phrases.iter().map(|s| s.as_str()).collect();
|
||||
|
||||
let embeddings = model.embed(texts, None)
|
||||
let embeddings = model.embedding.lock().embed(texts, None)
|
||||
.map_err(|e| format!("Embedding failed for '{}': {}", cmd.id, e))?;
|
||||
|
||||
// average all phrase vectors into one intent vector
|
||||
|
|
@ -166,9 +129,10 @@ fn build_intent_vectors(
|
|||
}
|
||||
|
||||
pub fn classify(text: &str) -> Result<(String, f64), String> {
|
||||
let mut classifier = CLASSIFIER.get().ok_or("Classifier not initialized")?.lock();
|
||||
let state = CLASSIFIER.get().ok_or("Classifier not initialized")?;
|
||||
|
||||
let embeddings = classifier.model.embed(vec![text], None)
|
||||
// only the embedding model needs locking, intents are read-only
|
||||
let embeddings = state.model.embedding.lock().embed(vec![text], None)
|
||||
.map_err(|e| format!("Failed to embed query: {}", e))?;
|
||||
|
||||
let mut query_vec = embeddings.into_iter().next()
|
||||
|
|
@ -182,11 +146,11 @@ pub fn classify(text: &str) -> Result<(String, f64), String> {
|
|||
}
|
||||
}
|
||||
|
||||
// cosine similarity against all intents (dot product of normalized vectors)
|
||||
let mut best_id = String::new();
|
||||
// cosine similarity - track index, clone only the winner
|
||||
let mut best_idx: usize = 0;
|
||||
let mut best_score: f64 = -1.0;
|
||||
|
||||
for intent in &classifier.intents {
|
||||
for (i, intent) in state.intents.iter().enumerate() {
|
||||
let score: f64 = query_vec.iter()
|
||||
.zip(intent.vector.iter())
|
||||
.map(|(a, b)| (*a as f64) * (*b as f64))
|
||||
|
|
@ -194,31 +158,16 @@ pub fn classify(text: &str) -> Result<(String, f64), String> {
|
|||
|
||||
if score > best_score {
|
||||
best_score = score;
|
||||
best_id = intent.id.clone();
|
||||
best_idx = i;
|
||||
}
|
||||
}
|
||||
|
||||
let best_id = state.intents[best_idx].id.clone();
|
||||
debug!("Embedding classify: '{}' -> '{}' ({:.2}%)", text, best_id, best_score * 100.0);
|
||||
|
||||
Ok((best_id, best_score))
|
||||
}
|
||||
|
||||
pub fn get_command<'a>(
|
||||
commands: &'a [JCommandsList],
|
||||
intent_id: &str,
|
||||
) -> Option<(&'a PathBuf, &'a crate::commands::JCommand)> {
|
||||
for cmd_list in commands {
|
||||
for cmd in &cmd_list.commands {
|
||||
if cmd.id == intent_id {
|
||||
return Some((&cmd_list.path, cmd));
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
// ### CACHE HELPERS
|
||||
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
struct CachedIntent {
|
||||
id: String,
|
||||
|
|
@ -243,4 +192,4 @@ fn load_cached_intents(path: &PathBuf) -> Result<Vec<IntentVector>, String> {
|
|||
id: c.id,
|
||||
vector: c.vector,
|
||||
}).collect())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,29 +1,27 @@
|
|||
use intent_classifier::{
|
||||
IntentClassifier, IntentPrediction, IntentError,
|
||||
IntentPrediction, IntentError,
|
||||
TrainingExample, TrainingSource, IntentId
|
||||
};
|
||||
|
||||
use tokio::sync::OnceCell;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::fs;
|
||||
|
||||
use crate::commands::{self, JCommand, JCommandsList};
|
||||
use crate::commands::{self, JCommandsList};
|
||||
use crate::models;
|
||||
use crate::models::intent_classifier::IntentClassifierModel;
|
||||
use crate::{APP_CONFIG_DIR, i18n};
|
||||
|
||||
static CLASSIFIER: OnceCell<IntentClassifier> = OnceCell::const_new();
|
||||
// static COMMANDS_MAP: OnceCell<Vec<JCommandsList>> = OnceCell::const_new();
|
||||
use once_cell::sync::OnceCell;
|
||||
|
||||
static MODEL: OnceCell<Arc<IntentClassifierModel>> = OnceCell::new();
|
||||
|
||||
const TRAINING_CACHE_FILE: &str = "intent_training.json";
|
||||
const COMMANDS_HASH_FILE: &str = "commands_hash.txt";
|
||||
|
||||
pub async fn init(commands: &[JCommandsList]) -> Result<(), String> {
|
||||
// parse commands first
|
||||
// let commands = commands::parse_commands()?;
|
||||
let current_hash = commands::commands_hash(&commands); // regen hash for current commands set
|
||||
let current_hash = commands::commands_hash(&commands);
|
||||
|
||||
// init classifier
|
||||
let classifier = IntentClassifier::new().await
|
||||
.map_err(|e| format!("Failed to init IntentClassifier: {}", e))?;
|
||||
let model = models::intent_classifier::load(models::registry(), "intent-classifier").await?;
|
||||
|
||||
// check if we can use cached training data
|
||||
let config_dir = APP_CONFIG_DIR.get().ok_or("Config dir not set")?;
|
||||
|
|
@ -39,10 +37,9 @@ pub async fn init(commands: &[JCommandsList]) -> Result<(), String> {
|
|||
|
||||
if should_retrain {
|
||||
info!("Training intent classifier with {} commands...", commands.len());
|
||||
train_classifier(&classifier, &commands).await?;
|
||||
train_classifier(&model.classifier, &commands).await?;
|
||||
|
||||
// save training data and hash
|
||||
if let Ok(export) = classifier.export_training_data().await {
|
||||
if let Ok(export) = model.classifier.export_training_data().await {
|
||||
let _ = fs::write(&cache_path, export);
|
||||
let _ = fs::write(&hash_path, ¤t_hash);
|
||||
info!("Training data cached.");
|
||||
|
|
@ -50,41 +47,23 @@ pub async fn init(commands: &[JCommandsList]) -> Result<(), String> {
|
|||
} else {
|
||||
info!("Loading cached training data...");
|
||||
if let Ok(data) = fs::read_to_string(&cache_path) {
|
||||
classifier.import_training_data(&data).await
|
||||
model.classifier.import_training_data(&data).await
|
||||
.map_err(|e| format!("Failed to import training data: {}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
// store data
|
||||
CLASSIFIER.set(classifier).map_err(|_| "Classifier already set")?;
|
||||
// COMMANDS_MAP.set(commands).map_err(|_| "Commands map already set")?;
|
||||
MODEL.set(model).map_err(|_| "Model already set")?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn classify(text: &str) -> Result<IntentPrediction, IntentError> {
|
||||
let classifier = CLASSIFIER.get().expect("IntentClassifier not initialized");
|
||||
classifier.predict_intent(text).await
|
||||
let model = MODEL.get().expect("IntentClassifier not initialized");
|
||||
model.classifier.predict_intent(text).await
|
||||
}
|
||||
|
||||
// get command by intent ID
|
||||
pub fn get_command(commands: &'static [JCommandsList], intent_id: &str) -> Option<(&'static PathBuf, &'static JCommand)> {
|
||||
// let commands = COMMANDS_MAP.get()?;
|
||||
|
||||
for assistant_cmd in commands {
|
||||
for cmd in &assistant_cmd.commands {
|
||||
if cmd.id == intent_id {
|
||||
return Some((&assistant_cmd.path, cmd));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
// based on: https://github.com/ciresnave/intent-classifier/blob/main/examples/basic_usage.rs
|
||||
async fn train_classifier(
|
||||
classifier: &IntentClassifier,
|
||||
classifier: &intent_classifier::IntentClassifier,
|
||||
commands: &[JCommandsList]
|
||||
) -> Result<(), String> {
|
||||
let lang = i18n::get_language();
|
||||
|
|
@ -94,7 +73,6 @@ async fn train_classifier(
|
|||
|
||||
for assistant_cmd in commands {
|
||||
for cmd in &assistant_cmd.commands {
|
||||
// use language-specific phrases
|
||||
let phrases = cmd.get_phrases(&lang);
|
||||
|
||||
for phrase in phrases.iter() {
|
||||
|
|
@ -115,4 +93,4 @@ async fn train_classifier(
|
|||
|
||||
info!("Added {} training examples for language '{}'", total_examples, lang);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue