Create app.py

app.py

@@ -0,0 +1,148 @@
+import gradio as gr
+import torch
+import os
+
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from langchain.chains import ConversationalRetrievalChain
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.memory import ConversationBufferMemory
+from langchain.llms import HuggingFaceHub
+from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
+from langchain.prompts import PromptTemplate
+#from langchain.chains import (
+#    StuffDocumentsChain, LLMChain, ConversationalRetrievalChain
+#)
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+# Static model name
+llm_name = "meta-llama/Llama-2-7b-chat-hf"
+
+# Static file path for multiple files
+static_file_paths = [
+    "IRM ISO_IEC_27001_2022(en) (1).pdf",
+    #"SCF - Cybersecurity & Data Privacy Capability Maturity Model (CP-CMM) (2023.4).pdf",
+    #"AG_Level1_V2.0_Final_20211210.pdf",
+    #"CIS_Controls_v8_v21.10.pdf",
+    #"CSF PDF v11.1.0-1.pdf",
+    #"ISO_31000_2018(en)-1.pdf",
+    #"OWASP Application Security Verification Standard 4.0.3-en-1.pdf",
+    #"NIST.CSWP.29.ipd The NIST Cybersecurity Framework 2.0 202308-1 (1).pdf",
+    #"ISO_IEC_27002_2022(en)-1.pdf",
+]
+
+# Use cuda for faster processing
+device = torch.device("cuda" if torch.cuda.is_available() else "CPU")
+
+# Load documents
+loaders = [PyPDFLoader(x) for x in static_file_paths]
+pages = []
+for loader in loaders:
+    pages.extend(loader.load())
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=600,
+    chunk_overlap=40,
+)
+doc_splits = text_splitter.split_documents(pages)
+embedding = HuggingFaceEmbeddings()
+vectordb = Chroma.from_documents(
+    documents=doc_splits,
+    embedding=embedding,
+)
+
+# Load model
+tokenizer = AutoTokenizer.from_pretrained(llm_name, token=os.environ['HUGGINGFACEHUB_API_TOKEN'],)
+model = AutoModelForCausalLM.from_pretrained(llm_name, token=os.environ['HUGGINGFACEHUB_API_TOKEN'], torch_dtype=torch.float16)
+model = model.to(device)
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512, device=device, token=os.environ['HUGGINGFACEHUB_API_TOKEN'])
+hf = HuggingFacePipeline(pipeline=pipe)
+
+# Set up template and memory
+template = """You are a helpful and appreciative cybersecurity expert who gives comprehensive answers using lists, step-by-step instructions and other aids. Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+{context}
+Question: {question}
+Helpful Answer:
+"""
+prompt = PromptTemplate.from_template(template)
+memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key='answer',
+        return_messages=True
+    )
+retriever = vectordb.as_retriever()
+qachain = ConversationalRetrievalChain.from_llm(
+    hf,
+    retriever=retriever,
+    chain_type="stuff", 
+    memory=memory,
+    return_source_documents=True,
+    combine_docs_chain_kwargs={
+        "prompt": prompt,
+    }
+)
+
+def format_chat_history(message, chat_history):
+    formatted_chat_history = []
+    for user_message, bot_message in chat_history:
+        formatted_chat_history.append(f"User: {user_message}")
+        formatted_chat_history.append(f"Assistant: {bot_message}")
+    return formatted_chat_history
+
+# Conversation with chatbot
+def conversation(qa_chain, message, history):
+    formatted_chat_history = format_chat_history(message, history)
+    response = qa_chain({"question": message, "chat_history": formatted_chat_history})
+    response_answer = response["answer"]
+    response_sources = response["source_documents"]
+    response_source1 = response_sources[0].page_content.strip()
+    response_source2 = response_sources[1].page_content.strip()
+    response_source1_page = response_sources[0].metadata["page"] + 1
+    response_source2_page = response_sources[1].metadata["page"] + 1
+    new_history = history + [(message, response_answer)]
+    return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page
+
+def demo():
+    with gr.Blocks(theme="base") as demo:
+        qa_chain = gr.State(qachain)
+
+        gr.Markdown(
+        """<center><h2>Context Chatbot</center></h2>
+        <h3>Ask any questions about your PDF documents, along with follow-ups</h3>
+        When generating answers, it takes past questions into account (via conversational memory), and includes document references for clarity purposes.</i>
+        """)
+
+        # Conversation with chatbot
+        with gr.Tab("Step 3 - Conversation with chatbot"):
+            chatbot = gr.Chatbot(height=600)
+            with gr.Row():
+                msg = gr.Textbox(placeholder="Type message", container=True)
+            with gr.Row():
+                submit_btn = gr.Button("Submit")
+                clear_btn = gr.ClearButton([msg, chatbot])
+            with gr.Accordion("Advanced - Document references", open=False):
+                with gr.Row():
+                    response_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
+                    response_source1_page = gr.Number(label="Page", scale=1)
+                with gr.Row():
+                    response_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
+                    response_source2_page = gr.Number(label="Page", scale=1)
+
+        # Preprocessing events
+        #db_btn.click(initialize_database, outputs=[vector_db, db_progress])
+
+        # Chatbot events
+        submit_btn.click(conversation, \
+            inputs=[qa_chain, msg, chatbot], \
+            outputs=[qa_chain, msg, chatbot, response_source1, response_source1_page, response_source2, response_source2_page], \
+            queue=False)
+        clear_btn.click(lambda:[None,"",0,"",0], \
+            inputs=None, \
+            outputs=[chatbot], \
+            queue=False)
+
+    demo.queue().launch(debug=True)
+
+if __name__ == "__main__":
+    demo()