feat: tantivy lexical index over Alpha-ID texts
This commit is contained in:
+74
-1
@@ -1 +1,74 @@
|
|||||||
// implemented in a later task
|
use crate::model::{AlphaIdEntry, AppError};
|
||||||
|
use crate::normalize::normalize;
|
||||||
|
use tantivy::collector::TopDocs;
|
||||||
|
use tantivy::query::QueryParser;
|
||||||
|
use tantivy::schema::{Schema, Value, STORED, TEXT};
|
||||||
|
use tantivy::{doc, Index, TantivyDocument};
|
||||||
|
|
||||||
|
pub struct LexicalIndex {
|
||||||
|
index: Index,
|
||||||
|
text_field: tantivy::schema::Field,
|
||||||
|
idx_field: tantivy::schema::Field,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl LexicalIndex {
|
||||||
|
fn schema() -> (Schema, tantivy::schema::Field, tantivy::schema::Field) {
|
||||||
|
let mut b = Schema::builder();
|
||||||
|
let text_field = b.add_text_field("text", TEXT);
|
||||||
|
let idx_field = b.add_u64_field("idx", STORED);
|
||||||
|
(b.build(), text_field, idx_field)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn build_in_ram(entries: &[AlphaIdEntry]) -> Result<Self, AppError> {
|
||||||
|
let (schema, text_field, idx_field) = Self::schema();
|
||||||
|
let index = Index::create_in_ram(schema);
|
||||||
|
let mut w = index
|
||||||
|
.writer(50_000_000)
|
||||||
|
.map_err(|e| AppError::Io(format!("tantivy writer: {e}")))?;
|
||||||
|
for (i, e) in entries.iter().enumerate() {
|
||||||
|
w.add_document(doc!(text_field => normalize(&e.text), idx_field => i as u64))
|
||||||
|
.map_err(|e| AppError::Io(format!("tantivy add: {e}")))?;
|
||||||
|
}
|
||||||
|
w.commit()
|
||||||
|
.map_err(|e| AppError::Io(format!("tantivy commit: {e}")))?;
|
||||||
|
Ok(Self {
|
||||||
|
index,
|
||||||
|
text_field,
|
||||||
|
idx_field,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Returns (entry_index, bm25_score) descending.
|
||||||
|
pub fn search(&self, query: &str, k: usize) -> Result<Vec<(usize, f32)>, AppError> {
|
||||||
|
let reader = self
|
||||||
|
.index
|
||||||
|
.reader()
|
||||||
|
.map_err(|e| AppError::Io(format!("tantivy reader: {e}")))?;
|
||||||
|
let searcher = reader.searcher();
|
||||||
|
let norm = normalize(query);
|
||||||
|
let sanitized = norm.replace(['(', ')', '"', ':'], " ");
|
||||||
|
// Use OR semantics (default) for recall-favoring behaviour; fall back to
|
||||||
|
// the raw normalised query on parse error.
|
||||||
|
let parsed = QueryParser::for_index(&self.index, vec![self.text_field])
|
||||||
|
.parse_query(&sanitized)
|
||||||
|
.or_else(|_| {
|
||||||
|
let mut qp = QueryParser::for_index(&self.index, vec![self.text_field]);
|
||||||
|
qp.set_conjunction_by_default();
|
||||||
|
qp.parse_query(&norm)
|
||||||
|
})
|
||||||
|
.map_err(|e| AppError::Parse(format!("query: {e}")))?;
|
||||||
|
let top = searcher
|
||||||
|
.search(&parsed, &TopDocs::with_limit(k.max(1)))
|
||||||
|
.map_err(|e| AppError::Io(format!("tantivy search: {e}")))?;
|
||||||
|
let mut out = Vec::new();
|
||||||
|
for (score, addr) in top {
|
||||||
|
let d: TantivyDocument = searcher
|
||||||
|
.doc(addr)
|
||||||
|
.map_err(|e| AppError::Io(format!("tantivy doc: {e}")))?;
|
||||||
|
if let Some(v) = d.get_first(self.idx_field).and_then(|v| v.as_u64()) {
|
||||||
|
out.push((v as usize, score));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Ok(out)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,23 @@
|
|||||||
|
use alpha_id::lexical::LexicalIndex;
|
||||||
|
use alpha_id::model::AlphaIdEntry;
|
||||||
|
|
||||||
|
fn entry(id: &str, code: &str, text: &str) -> AlphaIdEntry {
|
||||||
|
AlphaIdEntry {
|
||||||
|
alpha_id: id.into(), valid: true,
|
||||||
|
icd_primary: code.into(), icd_star: "".into(), icd_addon: "".into(),
|
||||||
|
icd_primary2: "".into(), orpha: "".into(), text: text.into(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn finds_entry_by_fuzzy_query() {
|
||||||
|
let entries = vec![
|
||||||
|
entry("I1", "E11.9", "Diabetes mellitus Typ 2"),
|
||||||
|
entry("I2", "I10.90", "Arterielle essentielle Hypertonie"),
|
||||||
|
];
|
||||||
|
let idx = LexicalIndex::build_in_ram(&entries).unwrap();
|
||||||
|
let hits = idx.search("diabetes typ 2", 5).unwrap();
|
||||||
|
assert!(!hits.is_empty());
|
||||||
|
assert_eq!(hits[0].0, 0); // entry index 0
|
||||||
|
assert!(hits[0].1 > 0.0); // bm25 score
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user