File size: 1,771 Bytes
970a8e9 fe3da19 970a8e9 fe3da19 970a8e9 fe3da19 5759a61 fe3da19 5759a61 fe3da19 | 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 | import os
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
os.environ["HF_HOME"] = "/tmp/hf_cache"
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
def load_pdf_document(file_path):
document_loader = PyPDFLoader(file_path)
return document_loader.load()
def chunk_documents(raw_documents):
text_processor = RecursiveCharacterTextSplitter(
chunk_size = 1000,
chunk_overlap = 200,
add_start_index = True
)
return text_processor.split_documents(raw_documents)
def find_related_documents(query, vector_database):
# return vector_database.similarity_search(query, k=2)
return vector_database.max_marginal_relevance_search(query, k=2, fetch_k=5, lambda_mult=0.6)
def ProcessDocuments(document_path: str) -> str:
loaded_doc = load_pdf_document(document_path)
chunked_doc = chunk_documents(loaded_doc)
vector_database = Chroma(
persist_directory=f"./chroma_store/{document_path.split('/')[-1].split('.')[0]}",
embedding_function=embedding_model
)
vector_database.add_documents(chunked_doc)
def generate_context(query: str, file: str):
ProcessDocuments(file)
vector_database = Chroma(
persist_directory=f"./chroma_store/{file.split('/')[-1].split('.')[0]}",
embedding_function=embedding_model
)
relevant_docs = find_related_documents(query, vector_database)
context_text = "\n".join([doc.page_content for doc in relevant_docs])
return query, context_text
|