RugMunchMedia/rmi-backend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
rmi-backend/app/splade_bm25.py

697 lines
22 KiB
Python
Raw Permalink Blame History

#!/usr/bin/env python3
"""
SPLADE + BM25 Sparse Search Engine for RMI RAG
================================================
Implements proper learned sparse representations (SPLADE-style) and BM25
with IDF weighting, replacing the naive token-overlap sparse search.
Architecture:
- SPLADE: Term-level importance weighting using the existing BGE-small + ReLU + log(1+abs())
- BM25: Standard Okapi BM25 with IDF, k1=1.2, b=0.75, stopword removal
- Integration: Drop-in replacement for _sparse_text_search in rag_service.py
The SPLADE approach generates sparse vectors where each dimension corresponds to a
vocabulary term, weighted by importance. This replaces naive token counting with
learned term expansion and importance weighting.
For production SPLADE, you'd use naver/splade-v3-doc or prithivida/Splade_PP_en_v2.
This module provides a lightweight approximation using BGE-small features + expansion,
with a clear upgrade path to a full SPLADE model.
"""
import asyncio
import logging
import math
import re
from collections import Counter
from typing import Optional
logger = logging.getLogger(__name__)
# ── Stopwords (crypto-domain-aware) ────────────────────────────────────
STOPWORDS = frozenset(
{
# Standard English
"a",
"an",
"the",
"is",
"are",
"was",
"were",
"be",
"been",
"being",
"have",
"has",
"had",
"do",
"does",
"did",
"will",
"would",
"could",
"should",
"may",
"might",
"must",
"shall",
"can",
"need",
"dare",
"ought",
"used",
"to",
"of",
"in",
"for",
"on",
"with",
"at",
"by",
"from",
"as",
"into",
"through",
"during",
"before",
"after",
"above",
"below",
"between",
"out",
"off",
"over",
"under",
"again",
"further",
"then",
"once",
"here",
"there",
"when",
"where",
"why",
"how",
"all",
"each",
"every",
"both",
"few",
"more",
"most",
"other",
"some",
"such",
"no",
"nor",
"not",
"only",
"own",
"same",
"so",
"than",
"too",
"very",
"just",
"because",
"but",
"and",
"or",
"if",
"while",
"although",
"this",
"that",
"these",
"those",
"i",
"me",
"my",
"myself",
"we",
"our",
"ours",
"ourselves",
"you",
"your",
"yours",
"yourself",
"he",
"him",
"his",
"himself",
"she",
"her",
"hers",
"herself",
"it",
"its",
"itself",
"they",
"them",
"their",
"theirs",
"themselves",
"what",
"which",
"who",
"whom",
"whose",
"about",
"up",
"down",
# Crypto-domain: common but uninformative
"token",
"coin",
"crypto",
"address",
"also",
"like",
"get",
"got",
"one",
"two",
"new",
"old",
"well",
"even",
"much",
"many",
"still",
}
)
# ── Crypto-Domain Term Expansion ───────────────────────────────────────
# Maps common crypto terms to related expansion terms for SPLADE-like expansion
CRYPTO_EXPANSIONS = {
"rug": {"rugpull", "rug pull", "scam", "exit scam", "honeypot", "dump"},
"rugpull": {"rug pull", "rug", "scam", "exit scam", "dump"},
"honeypot": {"honeypot", "scam", "rug", "trap", "cannot sell", "stuck funds"},
"whale": {"large holder", "whale wallet", "big mover", "institutional"},
"wash": {"wash trade", "wash trading", "volume manipulation", "fake volume"},
"sniper": {"sniper bot", "mev", "sandwich", "front-run", "frontrun"},
"pump": {"pump and dump", "pump", "artificial inflation", "manipulation"},
"dump": {"dump", "sell-off", "crash", "rug", "liquidation"},
"flash": {"flash loan", "flash attack", "instant loan", "defi exploit"},
"sybil": {"sybil attack", "fake accounts", "multi-account", "astroturf"},
"phishing": {"phishing", "scam", "fake site", "credential theft"},
"mixer": {"mixer", "tumbler", "privacy", "laundering", "tornado"},
"mint": {"minting", "token creation", "deploy", "launch"},
"bridge": {"bridge", "cross-chain", "interoperability"},
"stake": {"staking", "yield", "liquidity pool", "lp"},
"exploit": {"exploit", "vulnerability", "hack", "vuln", "attack"},
"sanction": {"ofac", "sanction", "blocked", "restricted"},
"frontier": {"frontier", "fibonacci", "retrace"},
}
# ── Simple Stemmer (Porter-lite for crypto terms) ─────────────────────
_STEM_RULES = [
(r"ing$", ""), # running → runn
(r"ed$", ""), # scammed → scamm
(r"ly$", ""), # previously → previous
(r"tion$", ""), # liquidation → liquidat
(r"ment$", ""), # manipulation → manipulat
(r"ness$", ""), # worthiness → worthi
(r"ity$", ""), # liquidity → liquid
(r"ous$", ""), # suspicious → suspici
(r"ive$", ""), # restrictive → restrict
]
def simple_stem(word: str) -> str:
"""Aggressive suffix stripping for search matching. Not proper linguistics."""
for pattern, replacement in _STEM_RULES:
new = re.sub(pattern, replacement, word)
if new != word and len(new) >= 3:
return new
return word
# ── Tokenizer ─────────────────────────────────────────────────────────
def tokenize(text: str, remove_stopwords: bool = True, stem: bool = True) -> list[str]:
"""
Tokenize text for BM25/SPLADE search:
1. Lowercase
2. Split on non-alphanumeric (preserves hex addresses like 0xabc...)
3. Remove stopwords
4. Simple stemming
"""
# Preserve hex addresses and dollar amounts
text = text.lower()
# Extract hex addresses, dollar amounts as single tokens
hex_pattern = r"(0x[a-f0-9]{6,})"
money_pattern = r"(\$[\d,.]+)"
tokens = []
# Keep hex addresses intact
hex_addrs = re.findall(hex_pattern, text)
money_vals = re.findall(money_pattern, text)
# Remove hex and money for general tokenization
clean = re.sub(hex_pattern, " ", text)
clean = re.sub(money_pattern, " ", clean)
clean = re.sub(r"[^\w\s-]", " ", clean)
raw_tokens = clean.split()
tokens.extend(hex_addrs)
tokens.extend(money_vals)
for tok in raw_tokens:
if len(tok) < 2:
continue
if remove_stopwords and tok in STOPWORDS:
continue
if stem:
tok = simple_stem(tok)
tokens.append(tok)
return tokens
# ── BM25 Implementation ────────────────────────────────────────────────
class BM25Index:
"""
Okapi BM25 with crypto-domain enhancements.
Parameters:
k1: Term frequency saturation (1.2-2.0 typical, lower = less TF saturation)
b: Document length normalization (0.75 typical, 1.0 = full normalization)
epsilon: Floor for IDF to avoid negative scores
"""
def __init__(self, k1: float = 1.2, b: float = 0.75, epsilon: float = 0.25):
self.k1 = k1
self.b = b
self.epsilon = epsilon
self.doc_count = 0
self.avg_dl = 0.0
self.doc_lengths: dict[str, int] = {} # doc_id -> length
self.doc_tokens: dict[str, Counter] = {} # doc_id -> token counts
self.df: Counter = Counter() # token -> document frequency
self._built = False
def add_document(self, doc_id: str, content: str, metadata: dict | None = None):
"""Add a document to the index."""
tokens = tokenize(content, remove_stopwords=True, stem=True)
tf = Counter(tokens)
self.doc_tokens[doc_id] = tf
self.doc_lengths[doc_id] = len(tokens)
self.doc_count += 1
# Update document frequencies
unique_tokens = set(tokens)
for token in unique_tokens:
self.df[token] += 1
# Store metadata if provided
if metadata and not hasattr(self, "_metadata"):
self._metadata = {}
if hasattr(self, "_metadata") and metadata:
self._metadata[doc_id] = metadata
# Store raw content for exact phrase matching in search_with_expansion
if not hasattr(self, "_raw_content"):
self._raw_content = {}
# Store a truncated version to limit memory usage
self._raw_content[doc_id] = content[:2000].lower()
def build(self):
"""Finalize the index (compute avg doc length, IDF values)."""
if self.doc_count == 0:
self.avg_dl = 1.0
else:
total_len = sum(self.doc_lengths.values())
self.avg_dl = total_len / self.doc_count if self.doc_count else 1.0
self._built = True
def save(self, path: str):
"""Persist BM25 index to disk via pickle."""
import os
import pickle
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "wb") as f:
pickle.dump(
{
"k1": self.k1,
"b": self.b,
"epsilon": self.epsilon,
"doc_count": self.doc_count,
"avg_dl": self.avg_dl,
"doc_lengths": self.doc_lengths,
"doc_tokens": dict(self.doc_tokens),
"df": dict(self.df),
"_built": self._built,
"_metadata": getattr(self, "_metadata", {}),
"_raw_content": getattr(self, "_raw_content", {}),
},
f,
)
logger.info(f"BM25 index persisted to {path} ({os.path.getsize(path) / 1024 / 1024:.1f}MB)")
@classmethod
def load(cls, path: str) -> Optional["BM25Index"]:
"""Load BM25 index from disk."""
import os
import pickle
if not os.path.exists(path):
return None
try:
with open(path, "rb") as f:
data = pickle.load(f)
idx = cls(k1=data["k1"], b=data["b"], epsilon=data["epsilon"])
idx.doc_count = data["doc_count"]
idx.avg_dl = data["avg_dl"]
idx.doc_lengths = data["doc_lengths"]
idx.doc_tokens = data["doc_tokens"]
idx.df = Counter(data["df"])
idx._built = data["_built"]
if data.get("_metadata"):
idx._metadata = data["_metadata"]
if data.get("_raw_content"):
idx._raw_content = data["_raw_content"]
logger.info(f"BM25 index loaded from disk: {idx.doc_count} docs, {len(idx.df)} terms")
return idx
except Exception as e:
logger.warning(f"Failed to load BM25 from disk: {e}")
return None
def idf(self, token: str) -> float:
"""Compute IDF for a token."""
df = self.df.get(token, 0)
if df == 0:
return 0.0
return math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1.0)
def score_document(self, doc_id: str, query_tokens: list[str]) -> float:
"""Compute BM25 score for a single document against query tokens."""
if doc_id not in self.doc_tokens:
return 0.0
doc_tf = self.doc_tokens[doc_id]
dl = self.doc_lengths[doc_id]
score = 0.0
for token in query_tokens:
tf = doc_tf.get(token, 0)
if tf == 0:
continue
token_idf = self.idf(token)
# BM25 formula
numerator = tf * (self.k1 + 1)
denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_dl)
score += token_idf * numerator / denominator
return score
def search(self, query: str, limit: int = 20, min_score: float = 0.01) -> list[dict]:
"""
Search the BM25 index.
Returns list of dicts with: id, score, content_preview, metadata
"""
if not self._built:
self.build()
query_tokens = tokenize(query, remove_stopwords=True, stem=True)
# Expand query with crypto-domain terms
expanded_tokens = list(query_tokens)
for token in query_tokens:
if token in CRYPTO_EXPANSIONS:
expanded_tokens.extend(tokenize(" ".join(CRYPTO_EXPANSIONS[token])))
# Score all documents
scored = []
for doc_id in self.doc_tokens:
score = self.score_document(doc_id, expanded_tokens)
if score >= min_score:
entry = {
"id": doc_id,
"bm25_score": round(score, 6),
"collection": doc_id.split(":")[1] if ":" in doc_id else "unknown",
}
if hasattr(self, "_metadata") and doc_id in self._metadata:
entry["metadata"] = self._metadata[doc_id]
scored.append(entry)
scored.sort(key=lambda x: x["bm25_score"], reverse=True)
return scored[:limit]
def search_with_expansion(self, query: str, limit: int = 20) -> list[dict]:
"""
SPLADE-inspired search: expand query terms, weight by importance.
In full SPLADE, a transformer generates a log(1+ReLU(weight)) sparse vector.
Our approximation:
1. Tokenize + expand with crypto-domain synonyms
2. Weight each term by its IDF (rare terms = important)
3. Add positional bonus for exact phrase matches
4. Weight expansion terms lower than original terms
"""
if not self._built:
self.build()
original_tokens = tokenize(query, remove_stopwords=True, stem=True)
# Build weighted query: original tokens weight=1.0, expansion weight=0.5
weighted_tokens: Counter = Counter()
for token in original_tokens:
weighted_tokens[token] += 1.0
for token in original_tokens:
if token in CRYPTO_EXPANSIONS:
for exp_token in tokenize(" ".join(CRYPTO_EXPANSIONS[token])):
weighted_tokens[exp_token] += 0.5
# Score documents using weighted query
scored = []
for doc_id in self.doc_tokens:
doc_tf = self.doc_tokens[doc_id]
dl = self.doc_lengths[doc_id]
score = 0.0
for token, weight in weighted_tokens.items():
tf = doc_tf.get(token, 0)
if tf == 0:
continue
token_idf = self.idf(token)
# BM25 term score
numerator = tf * (self.k1 + 1)
denominator = tf + self.k1 * (1 - self.b + self.b * dl / self.avg_dl)
bm25_term = token_idf * numerator / denominator
# SPLADE-inspired: max activation * IDF weight
# This approximates log(1 + ReLU(w_ij)) for each term
score += bm25_term * weight
# Exact phrase bonus (SPLADE positional interaction approximation)
query_lower = query.lower()
# Check raw content if we have it stored
if score > 0 and hasattr(self, "_raw_content") and doc_id in self._raw_content:
if query_lower in self._raw_content[doc_id].lower():
score *= 1.5
if score > 0.01:
entry = {
"id": doc_id,
"sparse_score": round(score, 6),
"collection": doc_id.split(":")[1] if ":" in doc_id else "unknown",
}
if hasattr(self, "_metadata") and doc_id in self._metadata:
entry["metadata"] = self._metadata[doc_id]
scored.append(entry)
scored.sort(key=lambda x: x["sparse_score"], reverse=True)
return scored[:limit]
# ── Async BM25 Builder (for Redis collections) ──────────────────────
async def build_bm25_from_redis(
collections: list[str] | None = None,
max_per_collection: int = 50000,
) -> BM25Index:
"""
Build a BM25 index from Redis-stored documents.
Reads content from rag:{collection}:{doc_id} keys.
"""
import json
import os
# Try to load .env from multiple possible locations (Docker vs dev)
from pathlib import Path
import redis.asyncio as aioredis
for env_path in ["/app/.env", str(Path(__file__).parent.parent.parent / ".env")]:
try:
from dotenv import load_dotenv
load_dotenv(env_path, override=True)
break
except Exception:
pass
if collections is None:
collections = [
"wallet_profiles",
"token_analysis",
"known_scams",
"scam_patterns",
"forensic_reports",
"market_intel",
"contract_audits",
"news_articles",
"transaction_patterns",
"general",
]
r = aioredis.Redis(
host=os.environ.get("REDIS_HOST", "rmi-redis"),
port=int(os.environ.get("REDIS_PORT", "6379")),
password=os.environ.get("REDIS_PASSWORD", ""),
decode_responses=True,
)
index = BM25Index(k1=1.2, b=0.75)
total_indexed = 0
for coll in collections:
idx_key = f"rag:idx:{coll}"
try:
doc_ids = await r.smembers(idx_key)
if not doc_ids:
continue
# Limit per collection for memory
sample = list(doc_ids)[:max_per_collection]
logger.info(f"BM25 indexing {coll}: {len(sample)} docs")
# Batch fetch
batch_size = 500
for i in range(0, len(sample), batch_size):
batch = sample[i : i + batch_size]
keys = [f"rag:{coll}:{did}" for did in batch]
pipe = r.pipeline()
for k in keys:
pipe.get(k)
raw_docs = await pipe.execute()
for did, data in zip(batch, raw_docs, strict=False):
if not data:
continue
try:
doc = json.loads(data)
content = doc.get("content", "") or ""
metadata = doc.get("metadata", {}) or {}
# Use collection:doc_id as the index key
doc_key = f"{coll}:{did}"
index.add_document(doc_key, content, metadata)
total_indexed += 1
except (json.JSONDecodeError, Exception):
continue
except Exception as e:
logger.warning(f"BM25 indexing for {coll} failed: {e}")
index.build()
await r.aclose()
logger.info(
f"BM25 index built: {total_indexed} docs, {len(index.df)} unique terms, avg doc length {index.avg_dl:.1f}"
)
return index
# ── In-memory BM25 cache ───────────────────────────────────────────────
_bm25_index: BM25Index | None = None
_bm25_built_at: float = 0.0
_BM25_TTL = 3600 # Rebuild every hour
_bm25_lock = asyncio.Lock()
async def get_bm25_index(force_rebuild: bool = False) -> BM25Index:
"""Get or build the BM25 index (persisted to disk, 1h TTL in memory)."""
global _bm25_index, _bm25_built_at
import os
import time
BM25_PATH = "/app/data/bm25_index.pkl"
now = time.time()
if _bm25_index and not force_rebuild and (now - _bm25_built_at) < _BM25_TTL:
return _bm25_index
async with _bm25_lock:
if _bm25_index and not force_rebuild and (time.time() - _bm25_built_at) < _BM25_TTL:
return _bm25_index
# Try loading from disk first
if os.path.exists(BM25_PATH) and not force_rebuild:
loaded = await asyncio.get_running_loop().run_in_executor(None, BM25Index.load, BM25_PATH)
if loaded and loaded.doc_count > 0:
_bm25_index = loaded
_bm25_built_at = time.time()
logger.info(
f"BM25 loaded from disk: {loaded.doc_count} docs (skipped {time.time() - now:.1f}s rebuild)"
)
return _bm25_index
# Build from Redis and persist
_bm25_index = await build_bm25_from_redis()
_bm25_built_at = time.time()
# Persist to disk in background
if _bm25_index.doc_count > 0:
asyncio.create_task(_save_bm25_async(BM25_PATH))
return _bm25_index
async def _save_bm25_async(path: str):
"""Save BM25 index to disk in executor thread."""
try:
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, _bm25_index.save, path)
except Exception as e:
logger.debug(f"BM25 disk save failed (non-critical): {e}")
async def bm25_search(
query: str,
collections: list[str] | None = None,
limit: int = 20,
min_score: float = 0.01,
) -> list[dict]:
"""
Search using BM25 with crypto-domain query expansion.
Drop-in enhancement for _sparse_text_search.
"""
index = await get_bm25_index()
results = index.search(query, limit=limit, min_score=min_score)
return results
async def splade_search(
query: str,
collections: list[str] | None = None,
limit: int = 20,
) -> list[dict]:
"""
SPLADE-inspired sparse search with term expansion and importance weighting.
Uses BM25 as backbone + crypto-domain expansion for SPLADE-like sparse vectors.
"""
index = await get_bm25_index()
results = index.search_with_expansion(query, limit=limit)
return results
Reference in a new issue View git blame Copy permalink