feat: add relation classifier (embedding + logreg) with training pipeline

Author: rohansx

Cargo.lock

@@ -23,6 +23,7 @@ orp = "0.9"
 composable = "0.9"
 ndarray = { workspace = true }
 tokenizers = "0.22"
+fastembed = { workspace = true }
 
 [dev-dependencies]
 tempfile = "3"

crates/ctxgraph-extract/Cargo.toml

@@ -4,6 +4,7 @@ pub mod ner;
 pub mod nli;
 pub mod pipeline;
 pub mod rel;
+pub mod relclf;
 pub mod relex;
 pub mod remap;
 pub mod schema;

crates/ctxgraph-extract/src/lib.rs

@@ -308,6 +308,23 @@ impl ModelManager {
         }
     }
 
+    /// Check for a locally available fine-tuned relation classifier model.
+    ///
+    /// Looks for `relation_classifier/model_int8.onnx` (or `model.onnx`) and
+    /// `relation_classifier/tokenizer.json` in the cache directory.
+    ///
+    /// Returns `Some((model_path, tokenizer_path))` if found, `None` otherwise.
+    pub fn find_relation_classifier(&self) -> Option<std::path::PathBuf> {
+        let base = self.cache_dir.join("relation_classifier");
+
+        [
+            base.join("model_int8.onnx"),
+            base.join("model.onnx"),
+        ]
+        .into_iter()
+        .find(|p| p.exists())
+    }
+
     /// Check for locally exported gliner-relex ONNX model.
     ///
     /// The relex model must be exported manually using:

crates/ctxgraph-extract/src/model_manager.rs

@@ -13,6 +13,7 @@ use orp::params::RuntimeParameters;
 use orp::pipeline::Pipeline;
 
 use crate::ner::ExtractedEntity;
+use crate::relclf::RelationClassifier;
 use crate::relex::RelexEngine;
 use crate::schema::ExtractionSchema;
 
@@ -42,6 +43,9 @@ pub enum RelEngine {
 /// Cached relex engine (loaded once per process).
 static RELEX_ENGINE: std::sync::OnceLock<Option<RelexEngine>> = std::sync::OnceLock::new();
 
+/// Cached relation classifier (loaded once per process).
+static RELCLF_ENGINE: std::sync::OnceLock<Option<RelationClassifier>> = std::sync::OnceLock::new();
+
 /// Model-based relation extraction using gline-rs.
 ///
 /// Requires `gliner-multitask-large-v0.5` ONNX model.
@@ -185,7 +189,7 @@ impl RelEngine {
                         .collect();
 
                     for rel in &result.relations {
-                        if rel.confidence < 0.7 {
+                        if rel.confidence < 0.80 {
                             continue;
                         }
 
@@ -210,6 +214,57 @@ impl RelEngine {
             }
         }
 
+        // Relation classifier (embedding-based) — auto-enabled if model found.
+        let relclf = RELCLF_ENGINE.get_or_init(|| {
+            let mgr = crate::model_manager::ModelManager::new().ok()?;
+            let model_path = mgr.find_relation_classifier()?;
+            RelationClassifier::new(&model_path).ok()
+        });
+
+        if let Some(classifier) = relclf {
+            // Use fastembed for generating embeddings (same model as ctxgraph-embed).
+            static EMBED_ENGINE: std::sync::OnceLock<
+                Option<fastembed::TextEmbedding>,
+            > = std::sync::OnceLock::new();
+            let embed = EMBED_ENGINE.get_or_init(|| {
+                fastembed::TextEmbedding::try_new(
+                    fastembed::InitOptions::new(fastembed::EmbeddingModel::AllMiniLML6V2),
+                )
+                .ok()
+            });
+
+            if let Some(embed_model) = embed {
+                let embed_fn = |text: &str| -> Result<Vec<f32>, RelError> {
+                    let mut vecs = embed_model
+                        .embed(vec![text], None)
+                        .map_err(|e| RelError::Inference(e.to_string()))?;
+                    vecs.pop()
+                        .ok_or_else(|| RelError::Inference("empty embedding".into()))
+                };
+
+                if let Ok(clf_relations) =
+                    classifier.classify_batch(text, entities, &embed_fn)
+                {
+                    let existing: std::collections::HashSet<(String, String)> = relations
+                        .iter()
+                        .map(|r| (r.head.clone(), r.tail.clone()))
+                        .collect();
+
+                    for rel in clf_relations {
+                        if rel.confidence < 0.76 {
+                            continue;
+                        }
+                        // Only add relations for pairs not already covered
+                        if !existing.contains(&(rel.head.clone(), rel.tail.clone()))
+                            && !existing.contains(&(rel.tail.clone(), rel.head.clone()))
+                        {
+                            relations.push(rel);
+                        }
+                    }
+                }
+            }
+        }
+
         Ok(relations)
     }
 }