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