File size: 5,576 Bytes
e88c82c |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 |
import os
import tempfile
import streamlit as st
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.retrievers import EnsembleRetriever
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
from langchain_community.llms import CTransformers
from langchain_community.retrievers import BM25Retriever
from langchain_community.vectorstores import Chroma
from streamlit_extras.add_vertical_space import add_vertical_space
@st.cache_resource(ttl="1h")
def get_retriever(pdf_files):
"""get retriever"""
docs = []
temp_dir = tempfile.TemporaryDirectory()
for pdf_file in pdf_files:
temp_pdf_file_path = os.path.join(temp_dir.name, pdf_file.name)
with open(temp_pdf_file_path, "wb") as f:
f.write(pdf_file.getvalue())
loader = PyPDFLoader(temp_pdf_file_path)
docs.extend(loader.load())
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
chunk_size=1500, chunk_overlap=200
)
chunks = text_splitter.split_documents(docs)
# get huggingface token from env secret
HF_TOKEN = os.environ.get("HF_TOKEN")
# embeddings
embeddings = HuggingFaceInferenceAPIEmbeddings(
api_key=HF_TOKEN,
model_name="BAAI/bge-base-en-v1.5",
)
# retrieve k
k = 5
# vector retriever
vector_store = Chroma.from_documents(chunks, embeddings)
vector_retriever = vector_store.as_retriever(search_kwargs={"k": k})
# semantic retriever
semantic_retriever = BM25Retriever.from_documents(chunks)
semantic_retriever.k = k
# ensemble retriever
ensemble_retriever = EnsembleRetriever(
retrievers=[vector_retriever, semantic_retriever], weights=[0.5, 0.5]
)
return ensemble_retriever
@st.cache_resource(ttl="1h")
def initialize_llm(_retriever):
"""initialize llm"""
# load llm model
model_type = "mistral"
model_id = "TheBloke/zephyr-7B-beta-GGUF"
model_file = "zephyr-7b-beta.Q4_K_S.gguf"
config = {
"max_new_tokens": 2048,
"repetition_penalty": 1.1,
"temperature": 1,
"top_k": 50,
"top_p": 0.9,
"stream": True,
"context_length": 4096,
"gpu_layers": 0,
"threads": int(os.cpu_count()),
}
llm = CTransformers(
model=model_id,
model_file=model_file,
model_type=model_type,
config=config,
lib="avx2",
)
chat_history = StreamlitChatMessageHistory()
# init chat history memory
memory = ConversationBufferMemory(
memory_key="chat_history", chat_memory=chat_history, return_messages=True
)
chain = ConversationalRetrievalChain.from_llm(
llm, retriever=_retriever, memory=memory, verbose=False
)
return chain, chat_history
def main():
"""main func"""
st.set_page_config(
page_title="Talk to PDF using Zephyr-7B-Beta",
page_icon="📰",
layout="centered",
initial_sidebar_state="expanded",
)
st.header("Talk to PDF files 📰", divider="rainbow")
st.subheader(
"Enjoy :red[talking] with :green[PDF] files using :sunglasses: Zephyr-7B-Beta"
)
st.markdown(
"""
* Used the [zephyr-7b-beta.Q4_K_S.gguf](https://huggingface.co/TheBloke/zephyr-7B-alpha-GGUF/blob/main/zephyr-7b-alpha.Q4_K_S.gguf) quantised
version of [Zephyr-7B Beta](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta) model
from the [TheBloke/zephyr-7B-beta-GGUF](https://huggingface.co/TheBloke/zephyr-7B-beta-GGUF) repositry.
___
"""
)
st.sidebar.title("Talk to PDF 📰")
st.sidebar.markdown(
"[Checkout the repository](https://github.com/ThivaV/chat_with_pdf_using_zephyr-7b)"
)
st.sidebar.markdown(
"""
### This is a LLM powered chatbot, built using:
* [Streamlit](https://streamlit.io)
* [LangChain](https://python.langchain.com/)
* [HuggingFaceH4/zephyr-7b-beta](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta)
* [TheBloke/zephyr-7B-beta-GGUF](https://huggingface.co/TheBloke/zephyr-7B-beta-GGUF)
* [CTransformers](https://github.com/marella/ctransformers)
* [Embeddings](https://huggingface.co/BAAI/bge-base-en-v1.5)
* [Chroma](https://docs.trychroma.com/?lang=py)
___
"""
)
add_vertical_space(2)
upload_pdf_files = st.sidebar.file_uploader(
"Upload a pdf files 📤", type="pdf", accept_multiple_files=True
)
if not upload_pdf_files:
st.info("👈 :red[Please upload pdf files] ⛔")
st.stop()
retriever = get_retriever(upload_pdf_files)
chain, chat_history = initialize_llm(retriever)
# load previous chat history
# re-draw the chat history in the chat window
for message in chat_history.messages:
st.chat_message(message.type).write(message.content)
if prompt := st.chat_input("Ask questions"):
with st.chat_message("human"):
st.markdown(prompt)
response = chain.invoke(prompt)
with st.chat_message("ai"):
st.write(response["answer"])
if __name__ == "__main__":
# init main func
main()
|