File size: 4,735 Bytes
1869fec
c6308ec
 
2f49f39
 
 
 
c6308ec
2f49f39
c6308ec
2f49f39
c6308ec
2f49f39
 
c6308ec
 
 
 
 
 
 
 
 
2f49f39
 
 
 
 
c6308ec
2f49f39
 
 
 
 
 
 
 
 
 
c6308ec
2f49f39
 
 
 
c6308ec
 
2f49f39
 
 
 
c6308ec
2f49f39
 
c6308ec
2f49f39
 
 
c6308ec
 
 
 
 
 
 
 
 
 
 
 
2f49f39
 
 
 
c6308ec
2f49f39
 
 
 
c6308ec
2af25b0
 
2f49f39
2af25b0
 
 
 
 
 
c6308ec
2af25b0
 
 
 
 
2f49f39
2af25b0
 
 
 
 
 
 
c6308ec
2f49f39
2af25b0
c6308ec
2f49f39
 
 
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
import streamlit as st
import os
from pathlib import Path
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 HuggingFaceEndpoint
from langchain.memory import ConversationBufferMemory
from unidecode import unidecode
import chromadb
import re

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"
]

def load_doc(list_file_path, chunk_size, chunk_overlap):
    loaders = [PyPDFLoader(x) for x in list_file_path]
    pages = []
    for loader in loaders:
        pages.extend(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):
    llm = HuggingFaceEndpoint(repo_id=llm_model, temperature=temperature, max_new_tokens=max_tokens, top_k=top_k)
    memory = ConversationBufferMemory(memory_key="chat_history", output_key='answer', return_messages=True)
    retriever = vector_db.as_retriever()
    qa_chain = ConversationalRetrievalChain.from_llm(
        llm,
        retriever=retriever,
        chain_type="stuff",
        memory=memory,
        return_source_documents=True,
        verbose=False
    )
    return qa_chain

def create_collection_name(file_path):
    collection_name = Path(file_path).stem
    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'
    return collection_name

def main():
    st.title("PDF-based Chatbot")

    uploaded_files = st.file_uploader("Upload PDF documents (single or multiple)", type="pdf", accept_multiple_files=True)

    if uploaded_files:
        chunk_size = st.slider("Chunk size", min_value=100, max_value=1000, value=600, step=20)
        chunk_overlap = st.slider("Chunk overlap", min_value=10, max_value=200, value=40, step=10)

        tabs = ["Process Document", "Initialize QA Chain", "Chatbot"]
        selected_tab = st.radio("Select Tab", tabs)

        if selected_tab == "Process Document":
            if st.button("Generate Vector Database"):
                list_file_path = [file.name for file in uploaded_files]
                collection_name = create_collection_name(list_file_path[0])
                doc_splits = load_doc(list_file_path, chunk_size, chunk_overlap)
                vector_db = create_db(doc_splits, collection_name)

        elif selected_tab == "Initialize QA Chain":
            if "vector_db" in st.session_state:
                temperature = st.slider("Temperature", min_value=0.01, max_value=1.0, value=0.7, step=0.1)
                max_tokens = st.slider("Max Tokens", min_value=224, max_value=4096, value=1024, step=32)
                top_k = st.slider("Top-K Samples", min_value=1, max_value=10, value=3, step=1)

                llm_model = st.selectbox("Choose LLM Model", list_llm)

                if st.button("Initialize QA Chain"):
                    qa_chain = initialize_llmchain(llm_model, temperature, max_tokens, top_k, st.session_state["vector_db"])

        elif selected_tab == "Chatbot":
            if "qa_chain" in st.session_state:
                message = st.text_input("Type your message")
                if st.button("Submit"):
                    response = st.session_state["qa_chain"]({"question": message, "chat_history": []})
                    st.write("Assistant:", response["answer"])

if __name__ == "__main__":
    main()