Spaces:
Running
Running
| import numpy as np | |
| import pandas as pd | |
| import faiss | |
| from pathlib import Path | |
| from sentence_transformers import SentenceTransformer, util | |
| import streamlit as st | |
| class SemanticSearch: | |
| def __init__(self, shard_dir="compressed_shards"): | |
| self.shard_dir = Path(shard_dir) | |
| self.shard_dir.mkdir(exist_ok=True, parents=True) | |
| self.model = None | |
| self.index_shards = [] | |
| def load_model(_self): | |
| return SentenceTransformer('all-MiniLM-L6-v2') | |
| def initialize_system(self): | |
| self.model = self.load_model() | |
| self._load_index_shards() | |
| def _load_index_shards(self): | |
| """Load FAISS shards directly from local directory""" | |
| for shard_path in sorted(self.shard_dir.glob("*.index")): | |
| self.index_shards.append(faiss.read_index(str(shard_path))) | |
| def search(self, query, top_k=5): | |
| """Search across all shards""" | |
| query_embedding = self.model.encode([query], convert_to_numpy=True) | |
| all_scores = [] | |
| all_indices = [] | |
| for shard_idx, index in enumerate(self.index_shards): | |
| distances, indices = index.search(query_embedding, top_k) | |
| # Convert local indices to global shard offsets | |
| global_indices = [ | |
| self._calculate_global_index(shard_idx, idx) | |
| for idx in indices[0] | |
| ] | |
| all_scores.extend(distances[0]) | |
| all_indices.extend(global_indices) | |
| return self._process_results(np.array(all_scores), np.array(all_indices), top_k) | |
| def _calculate_global_index(self, shard_idx, local_idx): | |
| """Convert shard-local index to global index""" | |
| # Implement your specific shard indexing logic here | |
| # Example: return f"{shard_idx}-{local_idx}" | |
| return local_idx # Simple version if using unique IDs | |
| def _process_results(self, distances, indices, top_k): | |
| """Format search results""" | |
| results = pd.DataFrame({ | |
| 'global_index': indices, | |
| 'similarity': 1 - (distances / 2) # L2 to cosine approximation | |
| }) | |
| return results.sort_values('similarity', ascending=False).head(top_k) | |
| def search_with_threshold(self, query, top_k=5, similarity_threshold=0.6): | |
| """Threshold-filtered search""" | |
| results = self.search(query, top_k*2) | |
| filtered = results[results['similarity'] > similarity_threshold].head(top_k) | |
| return filtered.reset_index(drop=True) |