Delete text_engine.py

text_engine.py

@@ -1,225 +0,0 @@
-# text_engine.py
-import os
-import pickle
-import logging
-from typing import List, Optional
-import numpy as np
-from sentence_transformers import SentenceTransformer
-import faiss
-from rank_bm25 import BM25Okapi
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-
-class Text_Search_Engine:
-    def __init__(
-        self,
-        base_folder: str = "vector_store",
-        model_name: str = "sentence-transformers/LaBSE",
-        index_type: str = "flat",
-    ):
-        self.base_folder = base_folder
-        self.embeddings_folder = os.path.join(base_folder, "embeddings")
-        self.docs_folder = os.path.join(base_folder, "documents")
-        os.makedirs(self.embeddings_folder, exist_ok=True)
-        os.makedirs(self.docs_folder, exist_ok=True)
-
-        self.model = SentenceTransformer(model_name)
-        self.index: Optional[faiss.Index] = None
-        self.rows: List[dict] = []
-        self.texts: List[str] = []
-        self.bm25: Optional[BM25Okapi] = None
-        self.index_type = index_type
-
-    # -------------------------
-    # Index creation utilities
-    # -------------------------
-    def _create_index(self, dimension: int, embeddings: np.ndarray):
-        if self.index_type == "flat":
-            self.index = faiss.IndexFlatL2(dimension)
-        elif self.index_type == "ivf":
-            nlist = max(1, min(256, len(embeddings) // 10))
-            quantizer = faiss.IndexFlatL2(dimension)
-            self.index = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_L2)
-            self.index.train(np.array(embeddings).astype("float32"))
-        elif self.index_type == "hnsw":
-            self.index = faiss.IndexHNSWFlat(dimension, 32)
-        else:
-            raise ValueError(f"Unsupported index type: {self.index_type}")
-
-    def _persist(self):
-        try:
-            if self.index is not None:
-                faiss.write_index(self.index, os.path.join(self.embeddings_folder, "multilingual.index"))
-            with open(os.path.join(self.docs_folder, "rows.pkl"), "wb") as f:
-                pickle.dump(self.rows, f)
-            logger.info("Persisted index and rows to disk.")
-        except Exception as e:
-            logger.exception("Failed to persist index/rows: %s", e)
-
-    # -------------------------
-    # Core operations
-    # -------------------------
-    def encode_store(self, rows: List[dict], texts: List[str]):
-        try:
-            embeddings = self.model.encode(texts, convert_to_numpy=True)
-            dimension = embeddings.shape[1]
-            self._create_index(dimension, embeddings)
-            self.index.add(np.array(embeddings).astype("float32"))
-
-            self.rows = rows
-            self.texts = texts
-            tokenized_corpus = [t.lower().split() for t in texts]
-            self.bm25 = BM25Okapi(tokenized_corpus)
-
-            self._persist()
-            logger.info("Index built with %d rows (index_type=%s).", len(rows), self.index_type)
-        except Exception as e:
-            logger.exception("Error in encode_store: %s", e)
-            raise
-
-    def load(self):
-        try:
-            index_path = os.path.join(self.embeddings_folder, "multilingual.index")
-            rows_path = os.path.join(self.docs_folder, "rows.pkl")
-            if os.path.exists(index_path) and os.path.exists(rows_path):
-                self.index = faiss.read_index(index_path)
-                with open(rows_path, "rb") as f:
-                    self.rows = pickle.load(f)
-                self.texts = [r["_search_text"] for r in self.rows]
-                tokenized_corpus = [t.lower().split() for t in self.texts]
-                self.bm25 = BM25Okapi(tokenized_corpus)
-                logger.info("Loaded index and %d rows from disk.", len(self.rows))
-            else:
-                logger.info("No persisted index/rows found.")
-        except Exception as e:
-            logger.exception("Error in load: %s", e)
-            raise
-
-    def add_rows(self, new_rows: List[dict], new_texts: List[str]):
-        try:
-            if not new_rows:
-                return
-
-            new_embeddings = self.model.encode(new_texts, convert_to_numpy=True).astype("float32")
-            if self.index is None:
-                self._create_index(new_embeddings.shape[1], new_embeddings)
-                self.index.add(new_embeddings)
-            else:
-                if isinstance(self.index, faiss.IndexIVFFlat) and not self.index.is_trained:
-                    combined = np.vstack([self.model.encode(self.texts, convert_to_numpy=True).astype("float32"), new_embeddings]) if self.texts else new_embeddings
-                    self.index.train(combined)
-                self.index.add(new_embeddings)
-
-            self.rows.extend(new_rows)
-            self.texts.extend(new_texts)
-            tokenized_corpus = [t.lower().split() for t in self.texts]
-            self.bm25 = BM25Okapi(tokenized_corpus)
-
-            self._persist()
-            logger.info("Added %d new rows. Total rows: %d", len(new_rows), len(self.rows))
-        except Exception as e:
-            logger.exception("Error in add_rows: %s", e)
-            raise
-
-    # -------------------------
-    # Search methods
-    # -------------------------
-    def search(self, query: str, top_k: int = 3):
-        try:
-            if self.index is None:
-                return []
-            query_emb = self.model.encode([query], convert_to_numpy=True).astype("float32")
-            k = min(top_k, len(self.rows))
-            distances, indices = self.index.search(query_emb, k=k)
-            results = [
-                {**self.rows[i], "distance": float(distances[0][j])}
-                for j, i in enumerate(indices[0])
-            ]
-            return sorted(results, key=lambda x: x["distance"])
-        except Exception as e:
-            logger.exception("Error in search: %s", e)
-            return []
-
-    def hybrid_search(self, query: str, top_k: int = 3, alpha: float = 0.5):
-        try:
-            if self.index is None or self.bm25 is None:
-                return []
-    
-            # 🔹 Step 1: Encode query
-            query_emb = self.model.encode([query], convert_to_numpy=True).astype("float32")
-    
-            # 🔹 Step 2: Retrieve top candidates (IMPORTANT)
-            retrieve_k = min(20, len(self.texts))  # candidate pool
-            distances, indices = self.index.search(query_emb, k=retrieve_k)
-    
-            candidate_ids = indices[0]
-    
-            # 🔹 Step 3: Semantic scores (convert distance → similarity)
-            sem_scores = {}
-            for j, i in enumerate(candidate_ids):
-                sim = 1 / (1 + distances[0][j])
-                sem_scores[i] = sim
-    
-            # 🔹 Step 4: BM25 scores (only for candidates)
-            tokenized_query = query.lower().split()
-            bm25_scores = self.bm25.get_scores(tokenized_query)
-    
-            lex_scores = {i: bm25_scores[i] for i in candidate_ids}
-    
-            # 🔹 Step 5: NORMALIZATION (CRITICAL)
-            def normalize(scores_dict):
-                vals = list(scores_dict.values())
-                if not vals:
-                    return scores_dict
-                min_v, max_v = min(vals), max(vals)
-                if max_v - min_v == 0:
-                    return {k: 0.0 for k in scores_dict}
-                return {k: (v - min_v) / (max_v - min_v) for k, v in scores_dict.items()}
-    
-            sem_scores = normalize(sem_scores)
-            lex_scores = normalize(lex_scores)
-    
-            # 🔹 Step 6: Combine scores
-            combined = []
-            for i in candidate_ids:
-                sem = sem_scores.get(i, 0.0)
-                lex = lex_scores.get(i, 0.0)
-                score = alpha * sem + (1 - alpha) * lex
-    
-                combined.append({**self.rows[i], "score": float(score)})
-    
-            # 🔹 Step 7: Sort and return
-            combined = sorted(combined, key=lambda x: x["score"], reverse=True)
-    
-            return combined[:top_k]
-    
-        except Exception as e:
-            logger.exception("Error in hybrid_search: %s", e)
-            return []
-
-    # -------------------------
-    # Utilities
-    # -------------------------
-    def clear_vdb(self):
-        try:
-            if self.index is not None:
-                try:
-                    self.index.reset()
-                except Exception:
-                    self.index = None
-            self.rows = []
-            self.texts = []
-            self.bm25 = None
-
-            index_path = os.path.join(self.embeddings_folder, "multilingual.index")
-            docs_path = os.path.join(self.docs_folder, "rows.pkl")
-            if os.path.exists(index_path):
-                os.remove(index_path)
-            if os.path.exists(docs_path):
-                os.remove(docs_path)
-            logger.info("Cleared vector DB and persisted files.")
-        except Exception as e:
-            logger.exception("Error in clear_vdb: %s", e)
-            raise