CoolFace
Apppublic

LightRT/text2sql_backend

sourceHugging Faceupdated 1mo agoView on Hugging Face
0likes
embedding.py30 linesDownload Raw Back to src
1from langchain_chroma import Chroma2from langchain_huggingface import HuggingFaceEmbeddings3import uuid4from src.scheme import create_scheme5import os6from dotenv import load_dotenv7import chromadb8 9load_dotenv()10 11COLLECTION_NAME = "Text2SQL"12 13chroma_client = chromadb.CloudClient(14    api_key=os.getenv("CHROMA_API_KEY"),15    tenant=os.getenv("CHROMA_TENANT"),16    database=os.getenv("CHROMA_DATABASE"),17)18 19embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")20vectorstore = Chroma(collection_name=COLLECTION_NAME,embedding_function=embedding_model,client=chroma_client)21 22def create_embeddings(connection_url : str , user_id : str) :23    docs = create_scheme(connection_url)24 25    for doc in docs :26        doc.metadata['user_id'] = user_id27 28    ids = [str(uuid.uuid4()) for _ in docs]29 30    vectorstore.add_documents(documents=docs , ids=ids)