Spaces:
Runtime error
Runtime error
File size: 1,217 Bytes
ec96972 eb87b3b ec96972 eb87b3b ec96972 eb87b3b ec96972 eb87b3b ec96972 eb87b3b |
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 |
from sentence_transformers import SentenceTransformer
import numpy as np
import time
from embedder import get_model
# Use the preloaded model from embedder instead of creating a new instance
def retrieve_chunks(index, texts, query, k=5):
start_time = time.time()
print(f"Retrieving chunks for query: '{query[:50]}...'")
# Time query embedding
embed_start = time.time()
model = get_model() # Use the preloaded model
query_vec = model.encode([query])
embed_time = time.time() - embed_start
print(f"Query embedding took: {embed_time:.3f} seconds")
# Time FAISS search
search_start = time.time()
distances, indices = index.search(np.array(query_vec), k)
search_time = time.time() - search_start
print(f"FAISS search took: {search_time:.3f} seconds")
# Time result processing
process_start = time.time()
results = [texts[i] for i in indices[0]]
process_time = time.time() - process_start
print(f"Result processing took: {process_time:.3f} seconds")
total_time = time.time() - start_time
print(f"Total chunk retrieval took: {total_time:.3f} seconds")
print(f"Retrieved {len(results)} chunks")
return results
|