app.py

from transformers import pipeline
from langchain.llms import HuggingFacePipeline
import torch
import bitsandbytes as bnb
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig, pipeline, BitsAndBytesConfig


from langchain.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.document_loaders import TextLoader
from langchain.document_loaders import UnstructuredExcelLoader
from langchain.embeddings import HuggingFaceInstructEmbeddings
from langchain.memory import ConversationBufferWindowMemory
from langchain.prompts import ChatPromptTemplate
from langchain.memory import ConversationBufferWindowMemory
import gradio as gr
from controller import Controller

# Loading Model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                  # Load model weights in 4-bit format
    bnb_4bit_compute_type=torch.float16 # To avoid slow inference as input type into Linear4bit is torch.float16
)

MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, device_map="auto", torch_dtype=torch.float16, quantization_config=bnb_config
)

generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
generation_config.max_new_tokens = 2000
generation_config.temperature = 0.7
generation_config.do_sample = True

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    return_full_text=True,
    generation_config=generation_config,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.eos_token_id,
)
zephyr_llm = HuggingFacePipeline(pipeline=pipe)

"""--------------------------------------------Starting UI part--------------------------------------------"""
# Configurations
persist_directory = "db"
chunk_size = 150
chunk_overlap = 0

class Retriever:
    def __init__(self):
        self.text_retriever = None
        self.vectordb = None
        self.embeddings = None
        self.memory = ConversationBufferWindowMemory(k=2, return_messages=True)

    def create_and_add_embeddings(self, file):
        os.makedirs("db", exist_ok=True) # Recheck this and understand reason of above

        self.embeddings = HuggingFaceInstructEmbeddings(model_name="BAAI/bge-base-en-v1.5",
                                                        model_kwargs={"device": "cuda"})

        loader = UnstructuredExcelLoader(file)
        documents = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        texts = text_splitter.split_documents(documents)

        self.vectordb = Chroma.from_documents(documents=texts,
                                 embedding=self.embeddings,
                                 persist_directory=persist_directory)
        
        self.text_retriever = self.vectordb.as_retriever(search_kwargs={"k": 3})


    def retrieve_text(self, query):
        prompt_zephyr = ChatPromptTemplate.from_messages([
            ("system", "You are an helpful and harmless AI Assistant who is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user."),
            ("human", "Context: {context}\n <|user|>\n {question}\n<|assistant|>\n"),
            ])

        qa = RetrievalQA.from_chain_type(
            llm=zephyr_llm,
            chain_type="stuff",
            retriever=self.text_retriever,
            return_source_documents=False,
            verbose=False,
            chain_type_kwargs={"prompt": prompt_zephyr},
            memory=self.memory,
        )

        response = qa.run(query)
        return response

class Controller:
    def __init__(self):
        self.retriever = None
        self.query = ""

    def embed_document(self, file):
        if file is not None:
            self.retriever = Retriever()
            self.retriever.create_and_add_embeddings(file.name)

    def retrieve(self, query):
        texts = self.retriever.retrieve_text(query)
        return texts


# Gradio Demo for trying out the Application
import os
from controller import Controller
import gradio as gr

os.environ["TOKENIZERS_PARALLELISM"] = "false"
colors = ["#64A087", "green", "black"]

CSS = """
#question input {
    font-size: 16px;
}
#app-title {
    width: 100%;
    margin: auto;
}
#url-textbox {
    padding: 0 !important;
}
#short-upload-box .w-full {
    min-height: 10rem !important;
}

#select-a-file {
    display: block;
    width: 100%;
}
#file-clear {
    padding-top: 2px !important;
    padding-bottom: 2px !important;
    padding-left: 8px !important;
    padding-right: 8px !important;
	margin-top: 10px;
}
.gradio-container .gr-button-primary {
    background: linear-gradient(180deg, #CDF9BE 0%, #AFF497 100%);
    border: 1px solid #B0DCCC;
    border-radius: 8px;
    color: #1B8700;
}
.gradio-container.dark button#submit-button {
    background: linear-gradient(180deg, #CDF9BE 0%, #AFF497 100%);
    border: 1px solid #B0DCCC;
    border-radius: 8px;
    color: #1B8700
}
table.gr-samples-table tr td {
    border: none;
    outline: none;
}
table.gr-samples-table tr td:first-of-type {
    width: 0%;
}
div#short-upload-box div.absolute {
    display: none !important;
}
gradio-app > div > div > div > div.w-full > div, .gradio-app > div > div > div > div.w-full > div {
    gap: 0px 2%;
}
gradio-app div div div div.w-full, .gradio-app div div div div.w-full {
    gap: 0px;
}
gradio-app h2, .gradio-app h2 {
    padding-top: 10px;
}
#answer {
    overflow-y: scroll;
    color: white;
    background: #666;
    border-color: #666;
    font-size: 20px;
    font-weight: bold;
}
#answer span {
    color: white;
}
#answer textarea {
    color:white;
    background: #777;
    border-color: #777;
    font-size: 18px;
}
#url-error input {
    color: red;
}
"""

controller = Controller()


def process_pdf(file):
    if file is not None:
        controller.embed_document(file)
    return (
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(visible=True),
        gr.update(visible=True),
    )


def respond(message, history):
    botmessage = controller.retrieve(message)
    history.append((message, botmessage))
    return "", history


def clear_everything():
    return (None, None, None)


with gr.Blocks(css=CSS, title="") as demo:
    gr.Markdown("# Marketing Email Generator ", elem_id="app-title")
    gr.Markdown("## Upload a CSV and ask your query!", elem_id="select-a-file")
    gr.Markdown(
        "Drop your file here 👇",
        elem_id="select-a-file",
    )
    with gr.Row():
        with gr.Column(scale=3):
            upload = gr.File(label="Upload PDF", type="file")
            with gr.Row():
                clear_button = gr.Button("Clear", variant="secondary")

        with gr.Column(scale=6):
            chatbot = gr.Chatbot()
            with gr.Row().style(equal_height=True):
                with gr.Column(scale=8):
                    question = gr.Textbox(
                        show_label=False,
                        placeholder="e.g. What is the document about?",
                        lines=1,
                        max_lines=1,
                    ).style(container=False)
                with gr.Column(scale=1, min_width=60):
                    submit_button = gr.Button(
                        "Send your Request 🤖", variant="primary", elem_id="submit-button"
                    )

    upload.change(
        fn=process_pdf,
        inputs=[upload],
        outputs=[
            question,
            clear_button,
            submit_button,
            chatbot,
        ],
        api_name="upload",
    )
    question.submit(respond, [question, chatbot], [question, chatbot])
    submit_button.click(respond, [question, chatbot], [question, chatbot])
    clear_button.click(
        fn=clear_everything,
        inputs=[],
        outputs=[upload, question, chatbot],
        api_name="clear",
    )

if __name__ == "__main__":
    demo.launch(enable_queue=False, debug=True, share=False)