LawBot / src /dev /test.py
Vishwanath77's picture
Upload 78 files
b82f276 verified
# # import faiss
# # import numpy as np
# # # Set the embedding dimension
# # dimension = 768
# # num_vectors = 10
# # # Create random embeddings
# # np.random.seed(42)
# # embeddings = np.random.random((num_vectors, dimension)).astype('float32')
# # # Create a FAISS index
# # index = faiss.IndexFlatL2(dimension)
# # index.add(embeddings)
# # # Create a random query vector
# # query_vector = np.random.random((1, dimension)).astype('float32')
# # # Search in the index
# # distances, indices = index.search(query_vector, 5)
# # # Print results
# # print("Search successful!")
# # print("Nearest Neighbors:", indices)
# # print("Distances:", distances)
# import os
# os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" # Temporary fix
# os.environ["FAISS_NO_OPENMP"] = "1" # Prevent FAISS from using OpenMP
# import faiss
# from numpy.linalg import norm
# index = faiss.read_index("faiss_test_index.bin")
# print("Index loaded successfully!")
# print("Number of vectors in the index:", index.ntotal)
# def vector_db_retriever(query_embeddings):
# query_embeddings = query_embeddings/norm(query_embeddings[0])
# distances, indices = index.search(query_embeddings, 5)
# return indices, distances
# # from transformers import AutoModel
# # from numpy.linalg import norm
# # cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b))
# # model = AutoModel.from_pretrained('../../volumes/models/jina-embeddings-v2-base-en/', trust_remote_code=True) # trust_remote_code is needed to use the encode method
# # # embeddings = model.encode(['How is the weather today?', 'What is the current weather like today?'])
# # # print(cos_sim(embeddings[0], embeddings[1]))
# # print(model)
# # query = "Application to Jammu and Kashmir"
# # query_embeddings = model.encode([query])
# # query_embeddings = query_embeddings/norm(query_embeddings[0])
# # print("generating query embeddings")
# # # Search in the index
# # distances, indices = index.search(query_embeddings, 5)
# # print("getting distance")
# # # Print results
# # print("Search successful!")
# # print("Nearest Neighbors:", indices)
# # print("Distances:", distances)
from pymilvus import MilvusClient
client = MilvusClient("./milvus_demo.db")