Spaces:
Sleeping
Sleeping
File size: 2,488 Bytes
2db8ee1 de88b8d 2db8ee1 de88b8d 2db8ee1 | 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 | import faiss
import numpy as np
import pickle
import os
class ChunkIndex:
def __init__(self):
self._model = None
# Global index for ALL chunks
self.index = None
self.chunks = [] # List of dicts: {"source": str, "text": str}
@property
def model(self):
if self._model is None:
from sentence_transformers import SentenceTransformer
self._model = SentenceTransformer("all-MiniLM-L6-v2")
return self._model
def add_chunks(self, source, chunks):
"""
chunks: list of strings (the text segments)
source: filename or identifier
"""
if not chunks:
return
# 1. Store metadata
for text in chunks:
self.chunks.append({
"source": source,
"text": text
})
# 2. Embed
embeddings = self.model.encode(chunks)
embeddings = np.array(embeddings)
# Ensure 2D
if embeddings.ndim == 1:
embeddings = embeddings.reshape(1, -1)
embeddings = embeddings.astype("float32")
faiss.normalize_L2(embeddings)
# 3. Add to FAISS index
if self.index is None:
self.index = faiss.IndexFlatIP(embeddings.shape[1])
self.index.add(embeddings)
def search(self, query, k=6):
if self.index is None or self.index.ntotal == 0:
return []
q_emb = self.model.encode([query])
q_emb = np.array(q_emb).astype("float32")
faiss.normalize_L2(q_emb)
k = min(k, self.index.ntotal)
scores, idxs = self.index.search(q_emb, k)
results = []
for i in idxs[0]:
if i < len(self.chunks):
item = self.chunks[i]
results.append({
"source": item["source"],
"content": item["text"]
})
return results
def save_local(self, folder_path):
os.makedirs(folder_path, exist_ok=True)
faiss.write_index(self.index, os.path.join(folder_path, "index.faiss"))
with open(os.path.join(folder_path, "chunks.pkl"), "wb") as f:
pickle.dump(self.chunks, f)
def load_local(self, folder_path):
self.index = faiss.read_index(os.path.join(folder_path, "index.faiss"))
with open(os.path.join(folder_path, "chunks.pkl"), "rb") as f:
self.chunks = pickle.load(f)
|