Create app.py

app.py

@@ -0,0 +1,275 @@
+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
+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)