Spaces:
Sleeping
Sleeping
File size: 4,437 Bytes
5a9839d |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.vectorstores import Qdrant
from langchain_openai.embeddings import OpenAIEmbeddings
import fitz
import io
import tiktoken
import requests
import os
from utilities.debugger import dprint
def tiktoken_len(text):
tokens = tiktoken.encoding_for_model("gpt-4o").encode(
text,
)
return len(tokens)
def download_document(state, url, file_name, download_folder):
file_path = os.path.join(download_folder, file_name)
if not os.path.exists(download_folder):
os.makedirs(download_folder)
if not os.path.exists(file_path):
print(f"Downloading {file_name} from {url}...")
response = requests.get(url)
if response.status_code == 200:
with open(file_path, 'wb') as f:
f.write(response.content)
else:
dprint(state, f"Failed to download document from {url}. Status code: {response.status_code}")
else:
dprint(state, f"{file_name} already exists locally.")
return file_path
def get_documents(state):
for url in state.document_urls:
dprint(state, f"Downloading and loading document from {url}...")
file_name = url.split("/")[-1]
file_path = download_document(state, url, file_name, state.download_folder)
loader = PyMuPDFLoader(file_path)
loaded_document = loader.load()
single_text_document = "\n".join([doc.page_content for doc in loaded_document])
#state.add_loaded_document(loaded_document) # Append the loaded documents to the list
#state.add_single_text_document(single_text_document)
dprint(state, f"Number of pages: {len(loaded_document)}")
# lets get titles and metadata
pdf = fitz.open(file_path)
metadata = pdf.metadata
title = metadata.get('title', 'Document 1')
#state.add_metadata(metadata)
#state.add_title(title)
document = {
"url": url,
"title": title,
"metadata": metadata,
"single_text_document": single_text_document,
}
state.add_document(document)
dprint(state, f"Title of Document: {title}")
dprint(state, f"Full metadata for Document 1: {metadata}")
pdf.close()
dprint(state, f"documents: {state.documents}")
def create_chunked_documents(state):
get_documents(state)
# file_path_1 = "data/Blueprint-for-an-AI-Bill-of-Rights.pdf"
# file_path_2 = "data/NIST.AI.600-1.pdf"
# loader = PyMuPDFLoader(file_path_1)
# documents_1 = loader.load()
# loader = PyMuPDFLoader(file_path_2)
# documents_2 = loader.load()
# print(f"Number of pages in 1: {len(documents_1)}")
# print(f"Number of pages in 2: {len(documents_2)}")
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=state.chunk_size,
chunk_overlap=state.chunk_overlap,
length_function = tiktoken_len,
)
combined_document_objects = []
dprint(state, "Chunking documents and creating document objects")
for document in state.documents:
dprint(state, f"processing documend: {document['title']}")
text = document["single_text_document"]
dprint(state, text)
title = document["title"]
chunks_document = text_splitter.split_text(text)
dprint(state, len(chunks_document))
document_objects = [Document(page_content=chunk, metadata={"source": title, "document_id": "doc1"}) for chunk in chunks_document]
dprint(state, f"Number of chunks for Document: {len(chunks_document)}")
combined_document_objects = combined_document_objects + document_objects
state.add_combined_document_objects(combined_document_objects)
def create_vector_store(state):
create_chunked_documents(state)
embedding_model = OpenAIEmbeddings(model=state.embedding_model)
qdrant_vectorstore = Qdrant.from_documents(
documents=state.combined_document_objects,
embedding=embedding_model,
location=":memory:"
)
qdrant_retriever = qdrant_vectorstore.as_retriever()
state.set_retriever(qdrant_retriever)
return qdrant_retriever |