Spaces:
Runtime error
Runtime error
File size: 5,186 Bytes
b2b3e33 ff800db d2dadf2 0de7101 d2dadf2 b2b3e33 d2dadf2 0de7101 b2b3e33 d2dadf2 0de7101 d2dadf2 b2b3e33 0de7101 b2b3e33 d2dadf2 b2b3e33 d2dadf2 0de7101 b2b3e33 d2dadf2 ff9a830 da3f35d ff9a830 0de7101 d2dadf2 0de7101 d2dadf2 0de7101 c447d13 d2dadf2 0de7101 d2dadf2 0de7101 d2dadf2 0de7101 b2b3e33 c447d13 b2b3e33 d2dadf2 ff9a830 dfdac3f 67858ce dfdac3f 67858ce dfdac3f 0de7101 67858ce dfdac3f 67858ce dfdac3f 67858ce dfdac3f 0de7101 c447d13 0de7101 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | import os
import json
from typing import List, Dict
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import linear_kernel
from rank_bm25 import BM25Okapi
import re
import numpy as np
CORPUS_DIR = "../corpus"
def _tokenize(text: str):
return re.findall(r"\b\w+\b", text.lower())
class SimpleSearchEngine:
def __init__(self, corpus_dir: str = CORPUS_DIR):
self.corpus_dir = corpus_dir
self.documents: List[Dict] = []
# TF-IDF
self.vectorizer = None
self.doc_tfidf = None
# BM25
self.bm25 = None
self._bm25_tokens = []
def _load_documents(self):
docs = []
doc_id = 0
if not os.path.exists(self.corpus_dir):
print(f"Warning: Corpus directory '{self.corpus_dir}' does not exist.")
return
for root, _, files in os.walk(self.corpus_dir):
for fname in files:
if not fname.endswith(".json"):
continue
path = os.path.join(root, fname)
try:
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
text = data.get("tf_idf_text", "").strip()
if not text:
continue
docs.append({
"id": doc_id,
"title": data.get("title", os.path.splitext(fname)[0]),
"url": data.get("url", ""),
"path": path,
"text": text,
})
doc_id += 1
except Exception as e:
print(f"Error reading {path}: {e}")
self.documents = docs
print(f"Loaded {len(self.documents)} documents from {self.corpus_dir}")
def build_index(self):
self._load_documents()
texts = [doc["text"] for doc in self.documents]
if not texts:
print("No documents found to index.")
return
# TF-IDF
self.vectorizer = TfidfVectorizer(token_pattern=r"\b\w+\b")
self.doc_tfidf = self.vectorizer.fit_transform(texts)
print("TF-IDF index built.")
# BM25
self._bm25_tokens = [_tokenize(t) for t in texts]
self.bm25 = BM25Okapi(self._bm25_tokens)
print("BM25 index built.")
def _minmax_normalize(self, scores: np.ndarray):
min_s = scores.min()
max_s = scores.max()
if max_s == min_s:
return np.zeros_like(scores)
return (scores - min_s) / (max_s - min_s)
def search(self, query: str, top_k: int = 5):
if self.vectorizer is None or self.doc_tfidf is None:
raise RuntimeError("Index not built. Call build_index() first.")
query_vec = self.vectorizer.transform([query])
scores = linear_kernel(query_vec, self.doc_tfidf).flatten()
ranked_idx = scores.argsort()[::-1][:top_k]
results = []
for idx in ranked_idx:
doc = self.documents[idx]
results.append({
"title": doc["title"],
"url": doc["url"],
"path": doc["path"],
"text": doc["text"],
"score": float(scores[idx])
})
return results
def search_bm25(self, query: str, top_k: int = 5):
if self.bm25 is None:
raise RuntimeError("BM25 index not built. Call build_index() first.")
scores = np.array(self.bm25.get_scores(_tokenize(query)))
ranked_idx = scores.argsort()[::-1][:top_k]
results = []
for idx in ranked_idx:
doc = self.documents[idx]
results.append({
"title": doc["title"],
"url": doc["url"],
"path": doc["path"],
"text": doc["text"],
"score": float(scores[idx])
})
return results
def search_hybrid(self, query: str, top_k: int = 5, alpha: float = 0.5):
if self.vectorizer is None or self.doc_tfidf is None or self.bm25 is None:
raise RuntimeError("Indexes not built. Call build_index() first.")
# TF-IDF scores
query_vec = self.vectorizer.transform([query])
tfidf_scores = linear_kernel(query_vec, self.doc_tfidf).flatten()
# BM25 scores
bm25_scores = np.array(
self.bm25.get_scores(_tokenize(query)),
dtype=float
)
# Normalize
tfidf_norm = self._minmax_normalize(tfidf_scores)
bm25_norm = self._minmax_normalize(bm25_scores)
# Combine
hybrid_scores = alpha * tfidf_norm + (1 - alpha) * bm25_norm
ranked_idx = hybrid_scores.argsort()[::-1][:top_k]
results = []
for idx in ranked_idx:
doc = self.documents[idx]
results.append({
"title": doc["title"],
"url": doc["url"],
"path": doc["path"],
"text": doc["text"],
"score": float(hybrid_scores[idx])
})
return results
|