| | |
| |
|
| | |
| | |
| | |
| | |
| |
|
| | |
| |
|
| | import os |
| | os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" |
| | os.environ["CUDA_VISIBLE_DEVICES"] = "2" |
| |
|
| | import faiss |
| | import numpy as np |
| | import torch |
| | from sentence_transformers import SentenceTransformer |
| |
|
| | |
| | model_id = "Qwen/Qwen3-Embedding-4B" |
| | index_path = "/home/mshahidul/readctrl/data/vector_db/qwen_em/full_wikipedia_index.faiss" |
| |
|
| | |
| | print("Loading Index...") |
| | index = faiss.read_index(index_path) |
| | print(f"Index loaded successfully.") |
| | print(f"Total vectors in index: {index.ntotal}") |
| | print(f"Vector dimension: {index.d}") |
| |
|
| | |
| | print("Loading model for query embedding...") |
| | model = SentenceTransformer( |
| | model_id, |
| | trust_remote_code=True, |
| | device="cuda", |
| | model_kwargs={"torch_dtype": torch.bfloat16} |
| | ) |
| |
|
| | |
| | query = "What is the capital of France?" |
| | |
| | query_vector = model.encode([query], convert_to_numpy=True).astype('float32') |
| |
|
| | k = 5 |
| | distances, indices = index.search(query_vector, k) |
| |
|
| | |
| | print("\n--- Search Results ---") |
| | print(f"Query: {query}") |
| | for i in range(k): |
| | print(f"Result {i+1}: Index ID {indices[0][i]}, Distance: {distances[0][i]:.4f}") |
| |
|
| | if indices[0][0] == -1: |
| | print("\nError: The search returned -1. This usually means the index is empty or improperly trained.") |
| | else: |
| | print("\nSuccess: The index returned valid neighbors!") |