audit_assistant / auditqa /engine /vectorstore.py
ppsingh's picture
Update auditqa/engine/vectorstore.py
42b1ea0 verified
raw
history blame
No virus
1.17 kB
from langchain_community.embeddings import HuggingFaceEmbeddings, HuggingFaceInferenceAPIEmbeddings
from langchain_community.vectorstores import Chroma, Qdrant
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
from dotenv import load_dotenv
import os
provider_retrieval_model = "HF"
embeddingmodel = "sentence-transformers/all-MiniLM-l6-v2"
load_dotenv()
HF_Token = os.environ.get("HF_TOKEN")
client_path = "./vectorstore/"
collection_name = "collection"
provider_retrieval_model = "HF"
def create_vectorstore(docs):
if provider_retrieval_model == "HF":
qdrantClient = QdrantClient(path=client_path, prefer_grpc=True)
embeddings = HuggingFaceInferenceAPIEmbeddings(
api_key=HF_Token, model_name=embeddingmodel
)
dim = 384
qdrantClient.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
)
vectorstore = Qdrant(
client=qdrantClient,
collection_name=collection_name,
embeddings=embeddings,
)
vectorstore.add_documents(docs)