multimodal / src /utils /ingest_text.py
NEXAS's picture
Update src/utils/ingest_text.py
cef1053 verified
raw
history blame
4.08 kB
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