File size: 3,291 Bytes
9c830d9 31794c3 bf905de |
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 |
import streamlit as st
from dotenv import load_dotenv
from PyPDF2 import PdfReader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
# from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from htmlTemplates import css, bot_template, user_template
from langchain.llms import HuggingFaceHub
import os
# from transformers import T5Tokenizer, T5ForConditionalGeneration
# from langchain.callbacks import get_openai_callback
hub_token = os.environ["HUGGINGFACE_HUB_TOKEN"]
def split_pdfs(pdf_docs):
"""Splits a list of PDF documents into smaller chunks.
Args:
pdf_docs: A list of PDF documents.
Returns:
A list of lists of PDF documents, where each sublist contains a smaller chunk of the original PDF documents.
"""
pdf_chunks = []
for pdf_doc in pdf_docs:
# Split the PDF document into pages.
pdf_reader = PdfReader(pdf_doc)
pdf_pages = pdf_reader.pages
# Split the PDF pages into chunks.
pdf_chunks.append([])
for pdf_page in pdf_pages:
# Add the PDF page to the current chunk.
pdf_chunks[-1].append(pdf_page)
# If the chunk is too large, start a new chunk.
if len(pdf_chunks[-1]) >= 10:
pdf_chunks.append([])
return pdf_chunks
def generate_response(pdf_chunks, llm_model):
"""Generates a response to a query using a list of PDF documents and an LLM model.
Args:
pdf_chunks: A list of lists of PDF documents, where each sublist contains a smaller chunk of the original PDF documents.
llm_model: An LLM model.
Returns:
A response to the query.
"""
# Generate a summary of each PDF chunk.
pdf_summaries = []
for pdf_chunk in pdf_chunks:
# Generate a summary of the PDF chunk.
pdf_summary = llm_model.generate(
prompt=f"Summarize the following text:\n{get_pdf_text(pdf_chunk)}",
max_new_tokens=100
)
# Add the summary to the list of summaries.
pdf_summaries.append(pdf_summary)
# Generate a response to the query using the summaries of the PDF chunks.
response = llm_model.generate(
prompt=f"Answer the following question using the following summaries:\n{get_text_chunks(pdf_summaries)}\n\nQuestion:",
max_new_tokens=200
)
return response
def main():
load_dotenv()
st.set_page_config(page_title="Chat with multiple PDFs", page_icon=":books:")
st.write(css, unsafe_allow_html=True)
# Load the LLM model.
llm_model = HuggingFaceHub(repo_id="mistralai/Mistral-7B-v0.1", huggingfacehub_api_token=hub_token)
if "conversation" not in st.session_state:
st.session_state.conversation = None
if "chat_history" not in st.session_state:
st.session_state.chat_history = None
st.header("Chat with multiple PDFs :books:")
user_question = st.text_input("Ask a question about your documents:")
# If the user asked a question, generate a response.
if user_question:
# Split the PDF documents into smaller chunks.
pdf_chunks = split_pdfs(st.session_state.pdf_docs)
# Generate a response to the query.
response = generate_response(pdf_chunks, llm_model)
st.write(response)
main()
|