| |
|
|
| |
| |
| |
| |
|
|
| |
|
|
| 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!") |