feat(oss): port v3 pipeline with hybrid search, entity extraction, and additive scoring (#4805)
Co-authored-by: Soumil Rathi <soumilrathi@gmail.com> Co-authored-by: Saket Aryan <saketaryan2002@gmail.com> Co-authored-by: chaithanyak42 <chaithanya.kumar42a@gmail.com> Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
soumil-rathi
co-authored by
Soumil Rathi
Saket Aryan
chaithanyak42
kartik-mem0
Soumil Rathi
Saket Aryan
chaithanyak42
kartik-mem0
parent
57f944e18a
commit
a488e19044
@@ -0,0 +1,357 @@
|
||||
"""
|
||||
Entity extraction from text using spaCy NLP.
|
||||
|
||||
Extracts four types of entities from text:
|
||||
- **Proper nouns**: Capitalized multi-word sequences (person names, places, brands)
|
||||
- **Quoted text**: Text in single or double quotes (titles, specific terms)
|
||||
- **Noun compounds**: Multi-word noun phrases with specific modifiers (e.g., "machine learning")
|
||||
- **Noun fallback**: Single nouns from circumstantial compound patterns
|
||||
|
||||
Public API:
|
||||
extract_entities(text: str) -> List[Tuple[str, str]]
|
||||
|
||||
Internal:
|
||||
_extract_entities_from_doc(doc) -> List[Tuple[str, str]]
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Words that are too generic to be useful as entity heads
|
||||
_GENERIC_HEADS = {
|
||||
"thing", "stuff", "way", "time", "experience", "situation", "case",
|
||||
"fact", "matter", "issue", "idea", "thought", "feeling", "place",
|
||||
"area", "part", "kind", "type", "sort", "lot", "bit", "day", "year",
|
||||
"week", "month", "moment", "instance", "example", "technique",
|
||||
"method", "approach", "process", "step", "tool", "result", "outcome",
|
||||
"goal", "task", "item", "topic", "scale", "size", "level", "degree",
|
||||
"amount", "number", "style", "look", "color", "colour", "shape",
|
||||
"form", "piece", "section", "side", "end", "edge", "surface", "point",
|
||||
}
|
||||
|
||||
# Modifiers that describe circumstance, not content
|
||||
_CIRCUMSTANTIAL_MODS = {
|
||||
"solo", "individual", "team", "group", "joint", "collaborative",
|
||||
"first", "last", "next", "previous", "final", "initial", "main", "side",
|
||||
}
|
||||
|
||||
# Adjectives too vague to make a compound entity specific
|
||||
_NON_SPECIFIC_ADJ = {
|
||||
"many", "few", "several", "some", "any", "all", "most", "more",
|
||||
"less", "much", "little", "enough", "various", "numerous", "multiple",
|
||||
"countless", "great", "good", "bad", "nice", "terrible", "awful",
|
||||
"awesome", "amazing", "wonderful", "horrible", "excellent", "poor",
|
||||
"best", "worst", "fine", "okay", "new", "old", "recent", "past",
|
||||
"future", "current", "previous", "next", "last", "first", "latest",
|
||||
"early", "late", "former", "modern", "ancient", "big", "small",
|
||||
"large", "tiny", "huge", "enormous", "long", "short", "tall", "high",
|
||||
"low", "wide", "narrow", "thick", "thin", "deep", "shallow",
|
||||
"similar", "different", "same", "other", "another", "such", "certain",
|
||||
"important", "main", "major", "minor", "key", "primary", "real",
|
||||
"actual", "true", "whole", "entire", "full", "complete", "total",
|
||||
"basic", "simple", "interesting", "boring", "exciting", "special",
|
||||
"particular", "general", "common", "unique", "rare", "typical",
|
||||
"usual", "normal", "regular", "possible", "likely", "potential",
|
||||
"available", "necessary", "only", "solo", "individual", "team",
|
||||
"group", "joint", "collaborative", "final", "initial", "side",
|
||||
}
|
||||
|
||||
# Generic tail words to strip from compound entities
|
||||
_GENERIC_ENDINGS = {
|
||||
"work", "works", "job", "jobs", "task", "tasks", "stuff", "things",
|
||||
"thing", "info", "information", "details", "data", "content",
|
||||
"material", "materials", "activities", "activity", "efforts", "effort",
|
||||
"options", "option", "choices", "choice", "results", "result",
|
||||
"output", "outputs", "products", "product", "items", "item",
|
||||
}
|
||||
|
||||
# Capitalized single words that are too generic to be proper nouns
|
||||
_GENERIC_CAPS = {
|
||||
"works", "items", "things", "stuff", "resources", "options", "tips",
|
||||
"ideas", "steps", "ways", "methods", "tools", "features", "benefits",
|
||||
"examples", "details", "notes", "instructions", "guidelines",
|
||||
"recommendations", "suggestions", "overview", "summary", "conclusion",
|
||||
"introduction", "pros", "cons", "advantages", "disadvantages",
|
||||
}
|
||||
|
||||
# Markdown/formatting markers to skip during extraction
|
||||
_FORMATTING_MARKERS = {"*", "-", "+", "\u2022", "\u2013", "\u2014", "#", "##", "###", "**", "__"}
|
||||
|
||||
|
||||
def _is_sentence_start(tokens: list, idx: int) -> bool:
|
||||
"""Check if a token is at the start of a sentence or after formatting."""
|
||||
if idx == 0:
|
||||
return True
|
||||
tok = tokens[idx]
|
||||
if tok.is_sent_start:
|
||||
return True
|
||||
prev = tokens[idx - 1].text
|
||||
return prev in ".!?:" or prev in _FORMATTING_MARKERS or "\n" in prev
|
||||
|
||||
|
||||
def _strip_generic_ending(toks: list) -> list:
|
||||
"""Remove generic trailing words from compound token sequences."""
|
||||
if len(toks) <= 1:
|
||||
return toks
|
||||
last = toks[-1].lemma_.lower() if hasattr(toks[-1], "lemma_") else toks[-1].lower()
|
||||
return toks[:-1] if last in _GENERIC_ENDINGS and len(toks) > 2 else toks
|
||||
|
||||
|
||||
def _lemmatize_compound(toks: list) -> str:
|
||||
"""Join compound tokens, lemmatizing nouns."""
|
||||
return " ".join(t.lemma_ if t.pos_ == "NOUN" else t.text for t in toks)
|
||||
|
||||
|
||||
def _has_artifacts(txt: str) -> bool:
|
||||
"""Check for formatting artifacts that indicate non-entity text."""
|
||||
return any(
|
||||
[
|
||||
"**" in txt or "__" in txt or ":*" in txt,
|
||||
re.search(r"\s\*\s|\s\*$|^\*\s", txt),
|
||||
" " in txt or "\n" in txt or "\t" in txt,
|
||||
len(txt) > 100,
|
||||
txt.startswith(("\u2022", "-", "+", "\u2013", "\u2014")),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def extract_entities(text: str) -> List[Tuple[str, str]]:
|
||||
"""Extract named entities, quoted text, and noun compounds from text.
|
||||
|
||||
This is the public API that accepts a string. It loads the spaCy model
|
||||
internally and delegates to _extract_entities_from_doc().
|
||||
|
||||
Args:
|
||||
text: Input text to extract entities from.
|
||||
|
||||
Returns:
|
||||
Deduplicated list of (entity_type, entity_text) tuples.
|
||||
Entity types: PROPER, QUOTED, COMPOUND, NOUN.
|
||||
Returns empty list if spaCy is unavailable.
|
||||
"""
|
||||
from mem0.utils.spacy_models import get_nlp_full
|
||||
|
||||
nlp = get_nlp_full()
|
||||
if nlp is None:
|
||||
return []
|
||||
|
||||
doc = nlp(text)
|
||||
return _extract_entities_from_doc(doc)
|
||||
|
||||
|
||||
def extract_entities_batch(texts: List[str], batch_size: int = 32) -> List[List[Tuple[str, str]]]:
|
||||
"""Extract entities from multiple texts using spaCy's nlp.pipe() for batched NER.
|
||||
|
||||
Uses spaCy's efficient batch processing pipeline instead of calling
|
||||
nlp() individually per text. Significantly faster for multiple texts.
|
||||
|
||||
Args:
|
||||
texts: List of input texts to extract entities from.
|
||||
batch_size: Number of texts to process in each spaCy batch.
|
||||
|
||||
Returns:
|
||||
List of entity lists, one per input text. Each entity list contains
|
||||
(entity_type, entity_text) tuples. Returns list of empty lists if
|
||||
spaCy is unavailable.
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
from mem0.utils.spacy_models import get_nlp_full
|
||||
|
||||
nlp = get_nlp_full()
|
||||
if nlp is None:
|
||||
return [[] for _ in texts]
|
||||
|
||||
results = []
|
||||
for doc in nlp.pipe(texts, batch_size=batch_size):
|
||||
results.append(_extract_entities_from_doc(doc))
|
||||
return results
|
||||
|
||||
|
||||
def _extract_entities_from_doc(doc) -> List[Tuple[str, str]]:
|
||||
"""Extract entities from a spaCy Doc object.
|
||||
|
||||
Ported from platform's shared.core.utils.entity_extraction.extract_entities().
|
||||
"""
|
||||
entities: List[Tuple[str, str]] = []
|
||||
text = doc.text
|
||||
tokens = list(doc)
|
||||
|
||||
# === PROPER NOUN SEQUENCES ===
|
||||
i = 0
|
||||
while i < len(tokens):
|
||||
tok = tokens[i]
|
||||
if tok.text in _FORMATTING_MARKERS:
|
||||
i += 1
|
||||
continue
|
||||
is_cap = tok.text and tok.text[0].isupper()
|
||||
is_label = i + 1 < len(tokens) and tokens[i + 1].text == ":"
|
||||
|
||||
if is_cap and not is_label and tok.pos_ in {"PROPN", "NOUN", "ADJ"}:
|
||||
seq = [(tok, i)]
|
||||
j = i + 1
|
||||
while j < len(tokens):
|
||||
t = tokens[j]
|
||||
if (t.text and t.text[0].isupper()) or t.text.lower() in {
|
||||
"'s", "of", "the", "in", "and", "for", "at", "is",
|
||||
}:
|
||||
seq.append((t, j))
|
||||
j += 1
|
||||
else:
|
||||
break
|
||||
# Strip trailing function words
|
||||
while seq and seq[-1][0].text.lower() in {"of", "the", "in", "and", "for", "at", "is", "'s"}:
|
||||
seq.pop()
|
||||
if seq:
|
||||
has_mid_cap = any(
|
||||
not _is_sentence_start(tokens, idx)
|
||||
for (t, idx) in seq
|
||||
if t.text[0].isupper() and t.text.lower() not in {"'s", "of", "the", "in", "and", "for", "at", "is"}
|
||||
)
|
||||
if has_mid_cap:
|
||||
phrase = "".join(t.text_with_ws for (t, idx) in seq).strip()
|
||||
if len(phrase) > 2:
|
||||
entities.append(("PROPER", phrase))
|
||||
i = j
|
||||
else:
|
||||
i += 1
|
||||
|
||||
# === QUOTED TEXT ===
|
||||
for m in re.finditer(r'"([^"]+)"', text):
|
||||
if len(m.group(1).strip()) > 2:
|
||||
entities.append(("QUOTED", m.group(1).strip()))
|
||||
for m in re.finditer(r"(?:^|[\s\(\[{,;])'([^']+)'(?=[\s\.,;:!?\)\]]|$)", text):
|
||||
if len(m.group(1).strip()) > 2:
|
||||
entities.append(("QUOTED", m.group(1).strip()))
|
||||
|
||||
# === NOUN-NOUN COMPOUNDS ===
|
||||
for chunk in doc.noun_chunks:
|
||||
chunk_tokens = list(chunk)
|
||||
split_indices: list = []
|
||||
poss_splits: list = []
|
||||
for idx, tok in enumerate(chunk_tokens):
|
||||
if tok.dep_ == "case" and tok.text in {"'s", "\u2019s", "'"}:
|
||||
split_indices.append(idx)
|
||||
poss_splits.append(idx)
|
||||
elif tok.pos_ == "PUNCT" and tok.text in {"'", '"', "\u2018", "\u2019", "\u201c", "\u201d"}:
|
||||
split_indices.append(idx)
|
||||
|
||||
if split_indices:
|
||||
groups: list = []
|
||||
prev = 0
|
||||
for split_idx in split_indices:
|
||||
if split_idx > prev:
|
||||
groups.append(chunk_tokens[prev:split_idx])
|
||||
if split_idx in poss_splits:
|
||||
next_split = next((s for s in split_indices if s > split_idx), None)
|
||||
owned = chunk_tokens[split_idx + 1: next_split if next_split else len(chunk_tokens)]
|
||||
if owned:
|
||||
first_content = next((t for t in owned if t.pos_ not in {"PUNCT", "PART"}), None)
|
||||
if not (first_content and first_content.text and first_content.text[0].isupper()):
|
||||
prev = next_split if next_split else len(chunk_tokens)
|
||||
continue
|
||||
prev = split_idx + 1
|
||||
if prev < len(chunk_tokens):
|
||||
groups.append(chunk_tokens[prev:])
|
||||
else:
|
||||
groups = [chunk_tokens]
|
||||
|
||||
for group in groups:
|
||||
if not group:
|
||||
continue
|
||||
head = next((t for t in reversed(group) if t.pos_ in {"NOUN", "PROPN"}), None)
|
||||
if not head:
|
||||
continue
|
||||
head_generic = head.lemma_.lower() in _GENERIC_HEADS
|
||||
content = [
|
||||
t
|
||||
for t in group
|
||||
if t.pos_ not in {"DET", "PRON", "PUNCT", "PART", "ADP", "SCONJ", "NUM"} and (t.pos_ == "ADJ" or not t.is_stop)
|
||||
]
|
||||
if not content:
|
||||
continue
|
||||
|
||||
compound_toks = [t for t in content if t.dep_ == "compound"]
|
||||
adj_toks = [t for t in content if t.pos_ == "ADJ" or t.dep_ == "amod"]
|
||||
has_spec_adj = any(t.lemma_.lower() not in _NON_SPECIFIC_ADJ for t in adj_toks)
|
||||
if head_generic and not has_spec_adj and not compound_toks:
|
||||
continue
|
||||
|
||||
if compound_toks:
|
||||
is_circ = any(t.lemma_.lower() in _CIRCUMSTANTIAL_MODS for t in compound_toks)
|
||||
if is_circ:
|
||||
val = head.lemma_ if head.pos_ == "NOUN" else head.text
|
||||
if len(val) > 2:
|
||||
entities.append(("NOUN", val))
|
||||
else:
|
||||
filtered = _strip_generic_ending(
|
||||
[t for t in content if not (t.pos_ == "ADJ" and t.lemma_.lower() in _NON_SPECIFIC_ADJ)]
|
||||
)
|
||||
if filtered:
|
||||
phrase = _lemmatize_compound(filtered)
|
||||
if len(phrase) > 3 and " " in phrase:
|
||||
entities.append(("COMPOUND", phrase))
|
||||
elif len(content) > 1 and has_spec_adj:
|
||||
filtered = _strip_generic_ending(
|
||||
[t for t in content if not ((t.pos_ == "ADJ" or t.dep_ == "amod") and t.lemma_.lower() in _NON_SPECIFIC_ADJ)]
|
||||
)
|
||||
if filtered:
|
||||
phrase = _lemmatize_compound(filtered)
|
||||
if len(phrase) > 3 and " " in phrase:
|
||||
entities.append(("COMPOUND", phrase))
|
||||
|
||||
# === FALLBACK: Mis-tagged VERB heads ===
|
||||
processed = {e[1].lower() for e in entities if e[0] == "COMPOUND"}
|
||||
generic_verb_heads = _GENERIC_HEADS | {"find", "buy", "purchase", "sale", "deal", "trip", "visit"}
|
||||
|
||||
def collect_compounds(head):
|
||||
return [t for t in doc if t.head == head and t.dep_ == "compound"]
|
||||
|
||||
for tok in doc:
|
||||
if tok.pos_ == "VERB" and tok.dep_ in {"pobj", "dobj", "nsubj"}:
|
||||
comps = sorted(collect_compounds(tok), key=lambda t: t.i)
|
||||
if comps:
|
||||
phrase_toks = comps if tok.lemma_.lower() in generic_verb_heads else comps + [tok]
|
||||
phrase = " ".join(t.text for t in phrase_toks)
|
||||
if phrase.lower() not in processed and len(phrase) > 3 and " " in phrase:
|
||||
entities.append(("COMPOUND", phrase))
|
||||
processed.add(phrase.lower())
|
||||
|
||||
# === DEDUPLICATION & CLEANUP ===
|
||||
seen: set = set()
|
||||
deduped = []
|
||||
for t, e in entities:
|
||||
k = e.lower().strip()
|
||||
if k not in seen and len(k) > 2:
|
||||
seen.add(k)
|
||||
deduped.append((t, e))
|
||||
|
||||
cleaned: List[Tuple[str, str]] = []
|
||||
for etype, etext in deduped:
|
||||
txt = re.sub(r"^\*+\s*|\s*\*+$", "", etext.strip())
|
||||
txt = re.sub(r"\s*:+$", "", txt)
|
||||
txt = re.sub(r"^\d+\s*\.\s*", "", txt)
|
||||
if not txt or len(txt) <= 2 or _has_artifacts(txt):
|
||||
continue
|
||||
if etype == "PROPER" and " " not in txt and txt.lower() in _GENERIC_CAPS:
|
||||
continue
|
||||
cleaned.append((etype, txt))
|
||||
|
||||
# Keep best type per entity (PROPER > COMPOUND > QUOTED > NOUN)
|
||||
type_pri = {"PROPER": 0, "COMPOUND": 1, "QUOTED": 2, "NOUN": 3, "VERB": 4}
|
||||
best: dict = {}
|
||||
for t, e in cleaned:
|
||||
k = e.lower()
|
||||
if k not in best or type_pri.get(t, 99) < type_pri.get(best[k][0], 99):
|
||||
best[k] = (t, e)
|
||||
deduped = list(best.values())
|
||||
|
||||
# Remove entities that are substrings of longer entities
|
||||
all_lower = [e[1].lower() for e in deduped]
|
||||
return [(t, e) for t, e in deduped if not any(e.lower() != o and e.lower() in o for o in all_lower)]
|
||||
@@ -209,30 +209,6 @@ class VectorStoreFactory:
|
||||
return instance
|
||||
|
||||
|
||||
class GraphStoreFactory:
|
||||
"""
|
||||
Factory for creating MemoryGraph instances for different graph store providers.
|
||||
Usage: GraphStoreFactory.create(provider_name, config)
|
||||
"""
|
||||
|
||||
provider_to_class = {
|
||||
"memgraph": "mem0.memory.memgraph_memory.MemoryGraph",
|
||||
"neptune": "mem0.graphs.neptune.neptunegraph.MemoryGraph",
|
||||
"neptunedb": "mem0.graphs.neptune.neptunedb.MemoryGraph",
|
||||
"kuzu": "mem0.memory.kuzu_memory.MemoryGraph",
|
||||
"apache_age": "mem0.memory.apache_age_memory.MemoryGraph",
|
||||
"default": "mem0.memory.graph_memory.MemoryGraph",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create(cls, provider_name, config):
|
||||
class_type = cls.provider_to_class.get(provider_name, cls.provider_to_class["default"])
|
||||
try:
|
||||
GraphClass = load_class(class_type)
|
||||
except (ImportError, AttributeError) as e:
|
||||
raise ImportError(f"Could not import MemoryGraph for provider '{provider_name}': {e}")
|
||||
return GraphClass(config)
|
||||
|
||||
|
||||
class RerankerFactory:
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
BM25 lemmatization for consistent keyword matching.
|
||||
|
||||
Uses spaCy's lemmatizer for better handling of:
|
||||
- Verb forms: attending/attends/attended -> attend
|
||||
- Comparatives/superlatives: older/oldest -> old
|
||||
- Plurals: memories -> memory
|
||||
- Avoids over-stemming: organization != organize
|
||||
|
||||
Also includes original -ing forms alongside lemmas to handle cases
|
||||
where spaCy's context-dependent lemmatization produces inconsistent
|
||||
results (e.g., "meeting" as noun vs verb -> different lemmas).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def lemmatize_for_bm25(text: str) -> str:
|
||||
"""Lemmatize text for BM25 matching.
|
||||
|
||||
Returns space-joined lemmas for full-text search. Falls back to
|
||||
the original text if spaCy is unavailable.
|
||||
"""
|
||||
from mem0.utils.spacy_models import get_nlp_lemma
|
||||
|
||||
nlp = get_nlp_lemma()
|
||||
if nlp is None:
|
||||
return text
|
||||
|
||||
doc = nlp(text.lower())
|
||||
tokens = []
|
||||
|
||||
for token in doc:
|
||||
if token.is_punct or token.is_stop:
|
||||
continue
|
||||
|
||||
lemma = token.lemma_
|
||||
if lemma.isalnum():
|
||||
tokens.append(lemma)
|
||||
|
||||
# Also add original if it ends in -ing and differs from lemma.
|
||||
# This handles noun/verb ambiguity (meeting/meet, attending/attend).
|
||||
if token.text.endswith("ing") and token.text != lemma and token.text.isalnum():
|
||||
tokens.append(token.text)
|
||||
|
||||
return " ".join(tokens)
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
Scoring utilities for hybrid retrieval.
|
||||
|
||||
Provides:
|
||||
- **BM25 normalization**: Sigmoid normalization of raw BM25 scores to [0, 1].
|
||||
- **BM25 parameter selection**: Query-length-adaptive sigmoid parameters.
|
||||
- **Additive scoring**: Combined scoring with semantic + BM25 + entity boost.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
def get_bm25_params(query: str, *, lemmatized: Optional[str] = None) -> tuple:
|
||||
"""Get BM25 sigmoid parameters based on query length.
|
||||
|
||||
Longer queries tend to have higher raw BM25 scores, so we adjust
|
||||
the sigmoid midpoint and steepness accordingly.
|
||||
|
||||
Returns:
|
||||
(midpoint, steepness) for sigmoid normalization.
|
||||
"""
|
||||
if lemmatized is None:
|
||||
from mem0.utils.lemmatization import lemmatize_for_bm25
|
||||
|
||||
lemmatized = lemmatize_for_bm25(query)
|
||||
num_terms = len(lemmatized.split()) if lemmatized else 1
|
||||
|
||||
if num_terms <= 3:
|
||||
return 5.0, 0.7
|
||||
elif num_terms <= 6:
|
||||
return 7.0, 0.6
|
||||
elif num_terms <= 9:
|
||||
return 9.0, 0.5
|
||||
elif num_terms <= 15:
|
||||
return 10.0, 0.5
|
||||
else:
|
||||
return 12.0, 0.5
|
||||
|
||||
|
||||
def normalize_bm25(raw_score: float, midpoint: float, steepness: float) -> float:
|
||||
"""Normalize BM25 score to [0, 1] using logistic sigmoid.
|
||||
|
||||
Args:
|
||||
raw_score: Raw BM25 score (unbounded, typically 0-20+).
|
||||
midpoint: Score at which sigmoid outputs 0.5.
|
||||
steepness: Controls how quickly sigmoid transitions.
|
||||
|
||||
Returns:
|
||||
Normalized score in range [0, 1].
|
||||
"""
|
||||
return 1.0 / (1.0 + math.exp(-steepness * (raw_score - midpoint)))
|
||||
|
||||
|
||||
ENTITY_BOOST_WEIGHT = 0.5
|
||||
|
||||
|
||||
def score_and_rank(
|
||||
semantic_results: List[Dict[str, Any]],
|
||||
bm25_scores: Dict[str, float],
|
||||
entity_boosts: Dict[str, float],
|
||||
threshold: float,
|
||||
top_k: int,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Score candidates additively and return top-k results.
|
||||
|
||||
For each candidate:
|
||||
semantic_score is taken from the result's score field.
|
||||
combined = (semantic + bm25 + entity_boost) / max_possible
|
||||
|
||||
Threshold gates the semantic score BEFORE combining -- candidates
|
||||
below the threshold are excluded even if BM25/entity would boost them.
|
||||
|
||||
The divisor adapts based on which signals are active:
|
||||
- Semantic only: max_possible = 1.0
|
||||
- Semantic + BM25: max_possible = 2.0
|
||||
- Semantic + BM25 + entity: max_possible = 2.5
|
||||
- Semantic + entity (no BM25): max_possible = 1.5
|
||||
|
||||
Returns:
|
||||
List of scored result dicts sorted by combined score descending.
|
||||
"""
|
||||
has_bm25 = bool(bm25_scores)
|
||||
has_entity = bool(entity_boosts)
|
||||
|
||||
max_possible = 1.0
|
||||
if has_bm25:
|
||||
max_possible += 1.0
|
||||
if has_entity:
|
||||
max_possible += ENTITY_BOOST_WEIGHT
|
||||
|
||||
scored: List[Dict[str, Any]] = []
|
||||
|
||||
for result in semantic_results:
|
||||
mem_id = result.get("id")
|
||||
if mem_id is None:
|
||||
continue
|
||||
|
||||
semantic_score = result.get("score", 0.0)
|
||||
if semantic_score < threshold:
|
||||
continue
|
||||
|
||||
mem_id_str = str(mem_id)
|
||||
bm25_score = bm25_scores.get(mem_id_str, 0.0)
|
||||
entity_boost = entity_boosts.get(mem_id_str, 0.0)
|
||||
|
||||
raw_combined = semantic_score + bm25_score + entity_boost
|
||||
combined = min(raw_combined / max_possible, 1.0)
|
||||
|
||||
scored.append(
|
||||
{
|
||||
"id": mem_id_str,
|
||||
"score": combined,
|
||||
"score_breakdown": {
|
||||
"semantic": semantic_score,
|
||||
"bm25": bm25_score,
|
||||
"entity_boost": entity_boost,
|
||||
},
|
||||
"payload": result.get("payload"),
|
||||
}
|
||||
)
|
||||
|
||||
scored.sort(key=lambda x: x["score"], reverse=True)
|
||||
return scored[:top_k]
|
||||
@@ -0,0 +1,91 @@
|
||||
"""
|
||||
Shared spaCy model loader.
|
||||
|
||||
Consolidates spaCy model loading into a single module so that
|
||||
entity_extraction and lemmatization share one instance instead of
|
||||
each loading their own copy from disk.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import threading
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_nlp_full = None
|
||||
_nlp_lemma = None
|
||||
_load_failed_full = False
|
||||
_load_failed_lemma = False
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def _ensure_model_available():
|
||||
"""Download en_core_web_sm if spaCy is installed but model is missing."""
|
||||
try:
|
||||
import spacy
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"spaCy is not installed. Install it with: pip install mem0ai[nlp]"
|
||||
)
|
||||
|
||||
if not spacy.util.is_package("en_core_web_sm"):
|
||||
logger.info("Downloading spaCy model en_core_web_sm...")
|
||||
try:
|
||||
from spacy.cli import download
|
||||
|
||||
download("en_core_web_sm")
|
||||
logger.info("spaCy model en_core_web_sm downloaded successfully")
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to download spaCy model en_core_web_sm: {e}. "
|
||||
"Please install manually: python -m spacy download en_core_web_sm"
|
||||
) from e
|
||||
|
||||
|
||||
def get_nlp_full():
|
||||
"""Return spaCy model with all pipelines (NER, tagger, etc.) for entity extraction."""
|
||||
global _nlp_full, _load_failed_full
|
||||
if _load_failed_full:
|
||||
return None
|
||||
if _nlp_full is not None:
|
||||
return _nlp_full
|
||||
with _lock:
|
||||
if _nlp_full is not None:
|
||||
return _nlp_full
|
||||
if _load_failed_full:
|
||||
return None
|
||||
try:
|
||||
_ensure_model_available()
|
||||
import spacy
|
||||
|
||||
_nlp_full = spacy.load("en_core_web_sm")
|
||||
logger.info("spaCy full model loaded")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load spaCy full model: {e}")
|
||||
_load_failed_full = True
|
||||
return None
|
||||
return _nlp_full
|
||||
|
||||
|
||||
def get_nlp_lemma():
|
||||
"""Return spaCy model with only lemmatizer for BM25 text processing."""
|
||||
global _nlp_lemma, _load_failed_lemma
|
||||
if _load_failed_lemma:
|
||||
return None
|
||||
if _nlp_lemma is not None:
|
||||
return _nlp_lemma
|
||||
with _lock:
|
||||
if _nlp_lemma is not None:
|
||||
return _nlp_lemma
|
||||
if _load_failed_lemma:
|
||||
return None
|
||||
try:
|
||||
_ensure_model_available()
|
||||
import spacy
|
||||
|
||||
_nlp_lemma = spacy.load("en_core_web_sm", disable=["ner", "parser"])
|
||||
logger.info("spaCy lemma model loaded")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load spaCy lemma model: {e}")
|
||||
_load_failed_lemma = True
|
||||
return None
|
||||
return _nlp_lemma
|
||||
Reference in New Issue
Block a user