Spaces:
Sleeping
Sleeping
from llama_parse import LlamaParse | |
from langchain_chroma import Chroma | |
from qdrant_client import QdrantClient | |
from langchain_community.vectorstores.qdrant import Qdrant | |
from langchain.text_splitter import RecursiveCharacterTextSplitter | |
from langchain_community.embeddings.fastembed import FastEmbedEmbeddings | |
from langchain_community.document_loaders.directory import DirectoryLoader | |
import os | |
from fastembed import TextEmbedding | |
from typing import List | |
import nest_asyncio | |
nest_asyncio.apply() | |
llamaparse_api_key = os.getenv("LLAMA_CLOUD_API_KEY") | |
#qdrant_url = os.getenv("QDRANT_URL ") | |
#qdrant_api_key = os.getenv("QDRANT_API_KEY") | |
groq_api_key = os.getenv("GROQ_API_KEY") | |
parsed_data_file = r".\data\parsed_data.pkl" | |
output_md = r".\data\output.md" | |
loki = r"data" | |
import pickle | |
# Define a function to load parsed data if available, or parse if not | |
def load_or_parse_data(loc): | |
data_file = parsed_data_file | |
if os.path.exists(data_file): | |
# Load the parsed data from the file | |
with open(data_file, "rb") as f: | |
parsed_data = pickle.load(f) | |
else: | |
# Perform the parsing step and store the result in llama_parse_documents | |
parsingInstructiontest10k = """The provided document is an user guide or a manual. | |
It contains many images and tables. | |
Try to be precise while answering the questions""" | |
parser = LlamaParse(api_key=llamaparse_api_key, result_type="markdown", parsing_instruction=parsingInstructiontest10k) # type: ignore | |
llama_parse_documents = parser.load_data(loc) | |
# Save the parsed data to a file | |
with open(data_file, "wb") as f: | |
pickle.dump(llama_parse_documents, f) | |
# Set the parsed data to the variable | |
parsed_data = llama_parse_documents | |
return parsed_data | |
# Create vector database | |
def create_vector_database(loc): | |
""" | |
Creates a vector database using document loaders and embeddings. | |
This function loads urls, | |
splits the loaded documents into chunks, transforms them into embeddings using OllamaEmbeddings, | |
and finally persists the embeddings into a Chroma vector database. | |
""" | |
# Call the function to either load or parse the data | |
llama_parse_documents = load_or_parse_data(loc) | |
#print(llama_parse_documents[1].text[:100]) | |
#with open('data/output.md', 'a') as f: # Open the file in append mode ('a') | |
# for doc in llama_parse_documents: | |
# f.write(doc.text + '\n') | |
with open(output_md,'a', encoding='utf-8') as f: # Open the file in append mode ('a') | |
for doc in llama_parse_documents: | |
f.write(doc.text + '\n') | |
loader = DirectoryLoader(loki, glob="**/*.md", show_progress=True) | |
documents = loader.load() | |
# Split loaded documents into chunks | |
text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=100) | |
print('data chunckex') | |
docs = text_splitter.split_documents(documents) | |
print(len(docs)) | |
#len(docs) | |
#docs[0] | |
# Initialize Embeddings | |
embeddings = FastEmbedEmbeddings() # type: ignore | |
#embeddings = TextEmbedding() | |
print('Vector DB started!') | |
# Create and persist a Chroma vector database from the chunked documents | |
qdrant = Qdrant.from_documents( | |
documents=docs, | |
embedding=embeddings, | |
path=r".\data\local_qdrant", | |
#url=qdrant_url, | |
collection_name="rag" | |
#api_key=qdrant_api_key | |
) | |
# save to disk | |
#db2 = Chroma.from_documents(docs, embeddings, persist_directory="./chroma_db") | |
#docs = db2.similarity_search(query) | |
# load from disk | |
#db3 = Chroma(persist_directory="./chroma_db", embedding_function=embeddings) | |
#query it | |
#query = "what is the agend of Financial Statements for 2022 ?" | |
#found_doc = qdrant.similarity_search(query, k=3) | |
#print(found_doc[0][:100]) | |
# | |
print('Vector DB created successfully !') | |
#query = "Switching between external devices connected to the TV" | |
#found_doc = qdrant.similarity_search(query, k=3) | |
#print(found_doc) | |
return qdrant |