|
|
|
import os |
|
import sys |
|
from timeit import default_timer as timer |
|
from typing import List |
|
|
|
from langchain.document_loaders import PyPDFDirectoryLoader |
|
from langchain.embeddings import HuggingFaceInstructEmbeddings |
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
from langchain.vectorstores.base import VectorStore |
|
from langchain.vectorstores.chroma import Chroma |
|
from langchain.vectorstores.faiss import FAISS |
|
|
|
from app_modules.init import app_init, get_device_types |
|
from app_modules.llm_summarize_chain import SummarizeChain |
|
|
|
|
|
def load_documents(source_pdfs_path, urls) -> List: |
|
loader = PyPDFDirectoryLoader(source_pdfs_path, silent_errors=True) |
|
documents = loader.load() |
|
if urls is not None and len(urls) > 0: |
|
for doc in documents: |
|
source = doc.metadata["source"] |
|
filename = source.split("/")[-1] |
|
for url in urls: |
|
if url.endswith(filename): |
|
doc.metadata["url"] = url |
|
break |
|
return documents |
|
|
|
|
|
def split_chunks(documents: List, chunk_size, chunk_overlap) -> List: |
|
text_splitter = RecursiveCharacterTextSplitter( |
|
chunk_size=chunk_size, chunk_overlap=chunk_overlap |
|
) |
|
return text_splitter.split_documents(documents) |
|
|
|
|
|
llm_loader = app_init(False)[0] |
|
|
|
source_pdfs_path = ( |
|
sys.argv[1] if len(sys.argv) > 1 else os.environ.get("SOURCE_PDFS_PATH") |
|
) |
|
chunk_size = os.environ.get("CHUNCK_SIZE") |
|
chunk_overlap = os.environ.get("CHUNK_OVERLAP") |
|
|
|
sources = load_documents(source_pdfs_path, None) |
|
|
|
print(f"Splitting {len(sources)} PDF pages in to chunks ...") |
|
|
|
chunks = split_chunks( |
|
sources, chunk_size=int(chunk_size), chunk_overlap=int(chunk_overlap) |
|
) |
|
|
|
print(f"Summarizing {len(chunks)} chunks ...") |
|
start = timer() |
|
|
|
summarize_chain = SummarizeChain(llm_loader) |
|
result = summarize_chain.call_chain( |
|
{"input_documents": chunks}, |
|
None, |
|
None, |
|
True, |
|
) |
|
|
|
end = timer() |
|
print(f"Completed in {end - start:.3f}s") |
|
|
|
print("\n\n***Summary:") |
|
print(result["output_text"]) |
|
|
