Spaces:
Sleeping
Sleeping
File size: 2,689 Bytes
ee8fb16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import logging
import os
from typing import List
import shutil
# from langchain_openai import OpenAIEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from local_loader import load_data_files
from splitter import split_documents
from dotenv import load_dotenv
from time import sleep
EMBED_DELAY = 0.02 # 20 milliseconds
# This is to get the Streamlit app to use less CPU while embedding documents into Chromadb.
class EmbeddingProxy:
def __init__(self, embedding):
self.embedding = embedding
def embed_documents(self, texts: List[str]) -> List[List[float]]:
sleep(EMBED_DELAY)
return self.embedding.embed_documents(texts)
def embed_query(self, text: str) -> List[float]:
sleep(EMBED_DELAY)
return self.embedding.embed_query(text)
# This happens all at once, not ideal for large datasets.
def create_vector_db(texts, embeddings=None, collection_name="chroma"):
if not texts:
logging.warning("Empty texts passed in to create vector database")
# Select embeddings
if not embeddings:
openai_api_key = os.environ["OPENAI_API_KEY"]
# embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key, model="text-embedding-3-small")
embeddings = HuggingFaceEmbeddings()
proxy_embeddings = EmbeddingProxy(embeddings)
persist_directory = os.path.join("store/", collection_name)
if os.path.exists(persist_directory):
shutil.rmtree(persist_directory)
db = Chroma(collection_name=collection_name,
embedding_function=proxy_embeddings,
persist_directory=persist_directory)
try:
db.add_documents(texts)
except Exception as e:
logging.error(f"Error adding documents to Chroma: {e}")
# You might want to handle the error more specifically here,
# such as retrying or returning an error indicator.
return db
def find_similar(vs, query):
docs = vs.similarity_search(query)
return docs
def main():
load_dotenv()
docs = load_data_files(data_dir="data") # Load data from your 'data' folder
texts = split_documents(docs)
vs = create_vector_db(texts)
# Use a relevant query from your financial domain
results = find_similar(vs, query="What are the fees for an Equity Ordinary Account?")
MAX_CHARS = 300
print("=== Results ===")
for i, text in enumerate(results):
content = text.page_content
n = max(content.find(' ', MAX_CHARS), MAX_CHARS)
content = text.page_content[:n]
print(f"Result {i + 1}:\n {content}\n")
if __name__ == "__main__":
main() |