import os import streamlit as st import re from tempfile import NamedTemporaryFile import time import pathlib #from PyPDF2 import PdfReader from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline from langchain_community.llms import LlamaCpp from langchain.prompts import PromptTemplate from langchain.chains import LLMChain from langchain.callbacks.manager import CallbackManager from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler from langchain_community.embeddings import HuggingFaceEmbeddings from langchain.chains import RetrievalQA from langchain_community.vectorstores import FAISS from langchain.chains.question_answering import load_qa_chain from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT from langchain_community.document_loaders import TextLoader from langchain_community.document_loaders import PyPDFLoader from langchain.memory import ConversationBufferWindowMemory from langchain.memory import ConversationBufferMemory from langchain.chains import ConversationalRetrievalChain from langchain.memory.chat_message_histories.streamlit import StreamlitChatMessageHistory from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.llms import HuggingFaceHub SECRET_TOKEN = os.getenv("HF_TOKEN") os.environ["HUGGINGFACEHUB_API_TOKEN"] = SECRET_TOKEN # sidebar contents with st.sidebar: st.title('DOC-QA DEMO ') st.markdown(''' ## About Detail this application: - LLM model: Phi-2-4bit - Hardware resource : Huggingface space 8 vCPU 32 GB ''') def split_docs(documents,chunk_size=1000): text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=200) sp_docs = text_splitter.split_documents(documents) return sp_docs @st.cache_resource def load_llama2_llamaCpp(): core_model_name = "phi-2.Q4_K_M.gguf" #n_gpu_layers = 32 n_batch = 512 callback_manager = CallbackManager([StreamingStdOutCallbackHandler()]) llm = LlamaCpp( model_path=core_model_name, #n_gpu_layers=n_gpu_layers, n_batch=n_batch, callback_manager=callback_manager, verbose=True,n_ctx = 4096, temperature = 0.1, max_tokens = 128 ) return llm def set_custom_prompt(): custom_prompt_template = """ Use the following pieces of information from context to answer the user's question. If you don't know the answer, don't try to make up an answer. Context : {context} Question : {question} Please answer the questions in a concise and straightforward manner. Helpful answer: """ prompt = PromptTemplate(template=custom_prompt_template, input_variables=['context', 'question', ]) return prompt @st.cache_resource def load_embeddings(): embeddings = HuggingFaceEmbeddings(model_name = "thenlper/gte-base", model_kwargs = {'device': 'cpu'}) return embeddings def main(): data = [] sp_docs_list = [] msgs = StreamlitChatMessageHistory(key="langchain_messages") print(msgs) if "messages" not in st.session_state: st.session_state.messages = [] repo_id = "mistralai/Mistral-7B-Instruct-v0.2" llm = HuggingFaceHub( repo_id=repo_id, model_kwargs={"temperature": 0.1, "max_length": 128}) # llm = load_llama2_llamaCpp() qa_prompt = set_custom_prompt() embeddings = load_embeddings() uploaded_file = st.file_uploader('Choose your .pdf file', type="pdf") if uploaded_file is not None : with NamedTemporaryFile(dir='PDF', suffix='.pdf', delete=False) as f: f.write(uploaded_file.getbuffer()) print(f.name) #filename = f.name loader = PyPDFLoader(f.name) pages = loader.load_and_split() data.extend(pages) #st.write(pages) f.close() os.unlink(f.name) os.path.exists(f.name) if len(data) > 0 : embeddings = load_embeddings() sp_docs = split_docs(documents = data) st.write(f"This document have {len(sp_docs)} chunks") sp_docs_list.extend(sp_docs) try: db = FAISS.from_documents(sp_docs_list, embeddings) memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, input_key="query", output_key="result") qa_chain = RetrievalQA.from_chain_type( llm = llm, chain_type = "stuff", retriever = db.as_retriever(search_kwargs = {'k':3}), return_source_documents = True, memory = memory, chain_type_kwargs = {"prompt":qa_prompt}) for message in st.session_state.messages: with st.chat_message(message["role"]): st.markdown(message["content"]) # Accept user input if query := st.chat_input("What is up?"): # Display user message in chat message container with st.chat_message("user"): st.markdown(query) # Add user message to chat history st.session_state.messages.append({"role": "user", "content": query}) start = time.time() response = qa_chain({'query': query}) with st.chat_message("assistant"): st.markdown(response['result']) end = time.time() st.write("Respone time:",int(end-start),"sec") print(response) # Add assistant response to chat history st.session_state.messages.append({"role": "assistant", "content": response['result']}) with st.expander("See the related documents"): for count, url in enumerate(response['source_documents']): st.write(str(count+1)+":", url) clear_button = st.button("Start new convo") if clear_button : st.session_state.messages = [] qa_chain.memory.chat_memory.clear() except: st.write("Plaese upload your pdf file.") if __name__ == '__main__': main()