|
|
import gradio as gr |
|
|
import os |
|
|
|
|
|
from langchain_community.document_loaders import PyPDFLoader |
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
|
from langchain_community.vectorstores import Chroma |
|
|
from langchain.chains import ConversationalRetrievalChain |
|
|
from langchain_community.embeddings import HuggingFaceEmbeddings |
|
|
from langchain_community.llms import HuggingFacePipeline |
|
|
from langchain.chains import ConversationChain |
|
|
from langchain.memory import ConversationBufferMemory |
|
|
from langchain_community.llms import HuggingFaceEndpoint |
|
|
|
|
|
from pathlib import Path |
|
|
import chromadb |
|
|
from unidecode import unidecode |
|
|
|
|
|
from transformers import AutoTokenizer |
|
|
import transformers |
|
|
import torch |
|
|
import tqdm |
|
|
import accelerate |
|
|
import re |
|
|
|
|
|
|
|
|
static_pdf_link = "https://huggingface.co/spaces/CCCDev/PDFChat/resolve/main/Data-privacy-policy.pdf" |
|
|
|
|
|
list_llm = ["mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1", |
|
|
"mistralai/Mistral-7B-Instruct-v0.1", "google/gemma-7b-it", "google/gemma-2b-it", |
|
|
"HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1", |
|
|
"meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", |
|
|
"TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct", "tiiuae/falcon-7b-instruct", |
|
|
"google/flan-t5-xxl"] |
|
|
list_llm_simple = [os.path.basename(llm) for llm in list_llm] |
|
|
|
|
|
|
|
|
|
|
|
def load_doc(file_path, chunk_size, chunk_overlap): |
|
|
loader = PyPDFLoader(file_path) |
|
|
pages = loader.load() |
|
|
text_splitter = RecursiveCharacterTextSplitter( |
|
|
chunk_size=chunk_size, |
|
|
chunk_overlap=chunk_overlap) |
|
|
doc_splits = text_splitter.split_documents(pages) |
|
|
return doc_splits |
|
|
|
|
|
|
|
|
|
|
|
def create_db(splits, collection_name): |
|
|
embedding = HuggingFaceEmbeddings() |
|
|
new_client = chromadb.EphemeralClient() |
|
|
vectordb = Chroma.from_documents( |
|
|
documents=splits, |
|
|
embedding=embedding, |
|
|
client=new_client, |
|
|
collection_name=collection_name, |
|
|
) |
|
|
return vectordb |
|
|
|
|
|
|
|
|
|
|
|
def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()): |
|
|
progress(0.1, desc="Initializing HF tokenizer...") |
|
|
|
|
|
progress(0.5, desc="Initializing HF Hub...") |
|
|
if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1": |
|
|
llm = HuggingFaceEndpoint( |
|
|
repo_id=llm_model, |
|
|
temperature=temperature, |
|
|
max_new_tokens=max_tokens, |
|
|
top_k=top_k, |
|
|
load_in_8bit=True, |
|
|
) |
|
|
elif llm_model in ["HuggingFaceH4/zephyr-7b-gemma-v0.1", "mosaicml/mpt-7b-instruct"]: |
|
|
raise gr.Error("LLM model is too large to be loaded automatically on free inference endpoint") |
|
|
elif llm_model == "microsoft/phi-2": |
|
|
llm = HuggingFaceEndpoint( |
|
|
repo_id=llm_model, |
|
|
temperature=temperature, |
|
|
max_new_tokens=max_tokens, |
|
|
top_k=top_k, |
|
|
trust_remote_code=True, |
|
|
torch_dtype="auto", |
|
|
) |
|
|
elif llm_model == "TinyLlama/TinyLlama-1.1B-Chat-v1.0": |
|
|
llm = HuggingFaceEndpoint( |
|
|
repo_id=llm_model, |
|
|
temperature=temperature, |
|
|
max_new_tokens=250, |
|
|
top_k=top_k, |
|
|
) |
|
|
elif llm_model == "meta-llama/Llama-2-7b-chat-hf": |
|
|
raise gr.Error("Llama-2-7b-chat-hf model requires a Pro subscription...") |
|
|
else: |
|
|
llm = HuggingFaceEndpoint( |
|
|
repo_id=llm_model, |
|
|
temperature=temperature, |
|
|
max_new_tokens=max_tokens, |
|
|
top_k=top_k, |
|
|
) |
|
|
|
|
|
progress(0.75, desc="Defining buffer memory...") |
|
|
memory = ConversationBufferMemory( |
|
|
memory_key="chat_history", |
|
|
output_key='answer', |
|
|
return_messages=True |
|
|
) |
|
|
retriever = vector_db.as_retriever() |
|
|
progress(0.8, desc="Defining retrieval chain...") |
|
|
qa_chain = ConversationalRetrievalChain.from_llm( |
|
|
llm, |
|
|
retriever=retriever, |
|
|
chain_type="stuff", |
|
|
memory=memory, |
|
|
return_source_documents=True, |
|
|
verbose=False, |
|
|
) |
|
|
progress(0.9, desc="Done!") |
|
|
return qa_chain |
|
|
|
|
|
|
|
|
|
|
|
def create_collection_name(filepath): |
|
|
collection_name = Path(filepath).stem |
|
|
collection_name = collection_name.replace(" ", "-") |
|
|
collection_name = unidecode(collection_name) |
|
|
collection_name = re.sub('[^A-Za-z0-9]+', '-', collection_name) |
|
|
collection_name = collection_name[:50] |
|
|
if len(collection_name) < 3: |
|
|
collection_name = collection_name + 'xyz' |
|
|
if not collection_name[0].isalnum(): |
|
|
collection_name = 'A' + collection_name[1:] |
|
|
if not collection_name[-1].isalnum(): |
|
|
collection_name = collection_name[:-1] + 'Z' |
|
|
print('Filepath: ', filepath) |
|
|
print('Collection name: ', collection_name) |
|
|
return collection_name |
|
|
|
|
|
|
|
|
|
|
|
def initialize_database(chunk_size, chunk_overlap, progress=gr.Progress()): |
|
|
file_path = static_pdf_link |
|
|
progress(0.1, desc="Creating collection name...") |
|
|
collection_name = create_collection_name(file_path) |
|
|
progress(0.25, desc="Loading document...") |
|
|
doc_splits = load_doc(file_path, chunk_size, chunk_overlap) |
|
|
progress(0.5, desc="Generating vector database...") |
|
|
vector_db = create_db(doc_splits, collection_name) |
|
|
progress(0.9, desc="Done!") |
|
|
return vector_db, collection_name, "Complete!" |
|
|
|
|
|
|
|
|
def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()): |
|
|
llm_name = list_llm[llm_option] |
|
|
print("llm_name: ", llm_name) |
|
|
qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress) |
|
|
return qa_chain, "Complete!" |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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"] |
|
|
if response_answer.find("Helpful Answer:") != -1: |
|
|
response_answer = response_answer.split("Helpful Answer:")[-1] |
|
|
response_sources = response["source_documents"] |
|
|
response_source1 = response_sources[0].page_content.strip() |
|
|
response_source2 = response_sources[1].page_content.strip() |
|
|
response_source3 = response_sources[2].page_content.strip() |
|
|
response_source1_page = response_sources[0].metadata["page"] + 1 |
|
|
response_source2_page = response_sources[1].metadata["page"] + 1 |
|
|
response_source3_page = response_sources[2].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, response_source3, response_source3_page |
|
|
|
|
|
|
|
|
def demo(): |
|
|
with gr.Blocks(theme="base") as demo: |
|
|
vector_db = gr.State() |
|
|
qa_chain = gr.State() |
|
|
collection_name = gr.State() |
|
|
|
|
|
gr.Markdown( |
|
|
"""<center><h2>PDF-based chatbot</center></h2> |
|
|
<h3>Ask any questions about your PDF documents</h3>""") |
|
|
gr.Markdown( |
|
|
"""<b>Note:</b> This AI assistant, using Langchain and open-source LLMs, performs retrieval-augmented generation (RAG) from your PDF documents. \ |
|
|
The user interface explicitely shows multiple steps to help understand the RAG workflow. |
|
|
This chatbot takes past questions into account when generating answers (via conversational memory), and includes document references for clarity purposes.<br> |
|
|
<br><b>Warning:</b> This space uses the free CPU Basic hardware from Hugging Face. Some steps and LLM models used below (free inference endpoints) can take some time to generate a reply. |
|
|
""") |
|
|
|
|
|
with gr.Tab("Step 2 - Process document"): |
|
|
with gr.Row(): |
|
|
db_btn = gr.Radio(["ChromaDB"], label="Vector database type", value="ChromaDB", type="index", |
|
|
info="Choose your vector database") |
|
|
with gr.Accordion("Advanced options - Document text splitter", open=False): |
|
|
with gr.Row(): |
|
|
chunk_size = gr.Slider(64, 4096, value=512, step=32, label="Text chunk size", |
|
|
info="Text length of each document chunk being embedded into the vector database. Default is 512.") |
|
|
chunk_overlap = gr.Slider(0, 1024, value=24, step=8, label="Text chunk overlap", |
|
|
info="Text overlap between each document chunk being embedded into the vector database. Default is 24.") |
|
|
|
|
|
initialize_db = gr.Button("Process document") |
|
|
|
|
|
with gr.Row(): |
|
|
output_db = gr.Textbox(label="Database initialization steps", placeholder="", show_label=False) |
|
|
with gr.Accordion("Vector database collection details", open=False): |
|
|
collection = gr.Textbox(label="Collection name", placeholder="", show_label=False) |
|
|
|
|
|
with gr.Tab("Step 3 - Initialize LLM"): |
|
|
with gr.Row(): |
|
|
llm_options = gr.Dropdown(list_llm_simple, label="Choose open-source LLM", |
|
|
value="Mistral-7B-Instruct-v0.2", |
|
|
info="Choose among the proposed open-source LLMs") |
|
|
with gr.Accordion("Advanced LLM options", open=False): |
|
|
with gr.Row(): |
|
|
llm_temperature = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="LLM temperature", |
|
|
info="LLM sampling temperature, in [0.01,1.0] range. Default is 0.1") |
|
|
llm_max_tokens = gr.Slider(32, 1024, value=512, step=16, label="Max tokens", |
|
|
info="Maximum number of new tokens to be generated, in [32,1024] range. Default is 512") |
|
|
llm_top_k = gr.Slider(1, 40, value=20, step=1, label="Top K", |
|
|
info="The number of highest probability vocabulary tokens to keep for top-k-filtering. Default is 20.") |
|
|
|
|
|
initialize_llm = gr.Button("Initialize LLM") |
|
|
|
|
|
with gr.Row(): |
|
|
output_llm = gr.Textbox(label="LLM initialization steps", placeholder="", show_label=False) |
|
|
|
|
|
with gr.Tab("Step 4 - Start chatting"): |
|
|
chatbot = gr.Chatbot(label="PDF chatbot", height=500) |
|
|
msg = gr.Textbox(label="Your question", placeholder="Type your question here...", show_label=False) |
|
|
clear = gr.Button("Clear chat") |
|
|
|
|
|
with gr.Accordion("Document sources (3)", open=False): |
|
|
gr.Markdown("Source 1") |
|
|
response_src1 = gr.Textbox(label="Source 1", placeholder="", show_label=False) |
|
|
response_src1_page = gr.Number(label="Page number", value=0, precision=0, interactive=False) |
|
|
gr.Markdown("Source 2") |
|
|
response_src2 = gr.Textbox(label="Source 2", placeholder="", show_label=False) |
|
|
response_src2_page = gr.Number(label="Page number", value=0, precision=0, interactive=False) |
|
|
gr.Markdown("Source 3") |
|
|
response_src3 = gr.Textbox(label="Source 3", placeholder="", show_label=False) |
|
|
response_src3_page = gr.Number(label="Page number", value=0, precision=0, interactive=False) |
|
|
|
|
|
initialize_db.click(initialize_database, |
|
|
inputs=[chunk_size, chunk_overlap], |
|
|
outputs=[vector_db, collection_name, output_db]) |
|
|
initialize_llm.click(initialize_LLM, |
|
|
inputs=[llm_options, llm_temperature, llm_max_tokens, llm_top_k, vector_db], |
|
|
outputs=[qa_chain, output_llm]) |
|
|
msg.submit(conversation, |
|
|
inputs=[qa_chain, msg, chatbot], |
|
|
outputs=[chatbot, msg, chatbot, response_src1, response_src1_page, response_src2, response_src2_page, |
|
|
response_src3, response_src3_page]) |
|
|
clear.click(lambda: None, None, chatbot, queue=False) |
|
|
clear.click(lambda: None, None, msg, queue=False) |
|
|
|
|
|
return demo.queue().launch(debug=True) |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo() |
|
|
|
