Spaces:
Starting
on
CPU Upgrade
Starting
on
CPU Upgrade
import glob | |
import os | |
from langchain.text_splitter import RecursiveCharacterTextSplitter, SentenceTransformersTokenTextSplitter | |
from transformers import AutoTokenizer | |
from torch import cuda | |
from langchain_community.document_loaders import PyMuPDFLoader | |
from langchain_community.embeddings import HuggingFaceEmbeddings, HuggingFaceInferenceAPIEmbeddings | |
from langchain_community.vectorstores import Qdrant | |
device = 'cuda' if cuda.is_available() else 'cpu' | |
#from dotenv import load_dotenv | |
#load_dotenv() | |
#HF_token = os.environ["HF_TOKEN"] | |
path_to_data = "./data/" | |
def process_pdf(): | |
files = {'MWTS2021':'./data/MWTS2021.pdf', | |
'MWTS2022':'./data/MWTS2022.pdf', | |
'Consolidated2021':'./data/Consolidated2021.pdf'} | |
docs = {} | |
for file,value in files.items(): | |
try: | |
docs[file] = PyMuPDFLoader(value).load() | |
except Exception as e: | |
print("Exception: ", e) | |
# text splitter based on the tokenizer of a model of your choosing | |
# to make texts fit exactly a transformer's context window size | |
# langchain text splitters: https://python.langchain.com/docs/modules/data_connection/document_transformers/ | |
chunk_size = 256 | |
text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer( | |
AutoTokenizer.from_pretrained("BAAI/bge-small-en-v1.5"), | |
chunk_size=chunk_size, | |
chunk_overlap=10, | |
add_start_index=True, | |
strip_whitespace=True, | |
separators=["\n\n", "\n"], | |
) | |
all_documents = {'Consolidated':[], 'MWTS':[]} | |
for file,value in docs.items(): | |
doc_processed = text_splitter.split_documents(value) | |
for doc in doc_processed: | |
doc.metadata["source"] = file | |
doc.metadata["year"] = file[-4:] | |
for key in all_documents: | |
if key in file: | |
print(key) | |
all_documents[key].append(doc_processed) | |
for key, docs_processed in all_documents.items(): | |
docs_processed = [item for sublist in docs_processed for item in sublist] | |
all_documents[key] = docs_processed | |
embeddings = HuggingFaceEmbeddings( | |
model_kwargs = {'device': device}, | |
encode_kwargs = {'normalize_embeddings': True}, | |
model_name="BAAI/bge-small-en-v1.5" | |
) | |
qdrant_collections = {} | |
for file,value in all_documents.items(): | |
print("emebddings for:",file) | |
qdrant_collections[file] = Qdrant.from_documents( | |
value, | |
embeddings, | |
location=":memory:", | |
collection_name=file, | |
) | |
print("done") | |
return qdrant_collections |