Spaces:
Runtime error
Runtime error
| 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 | |