Update app.py

app.py

@@ -3,8 +3,40 @@ import time
 import streamlit as st
 from streamlit_chat import message
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import textwrap
 
-from chat import generate_response
+from chat import generate_response, generate_tag
+
+@st.cache(allow_output_mutation=True, hash_funcs={"_json.Scanner": lambda x: None})
+def create_database():
+    import json
+    from langchain.docstore.document import Document
+    from langchain.embeddings import HuggingFaceEmbeddings
+    from langchain.vectorstores import FAISS
+
+    # Maybe it is better to save the embeddings than the text dataset
+    if "db" not in st.session_state:
+        json_file_path = "./new_dataset.json"
+
+        string_chunks = []
+
+        with open(json_file_path, "r") as json_file:
+            for line in json_file:
+                if line != '\n':
+                    json_string = json.loads(line)
+                    string_chunks.append(json_string)
+        documents_ = []
+        for line in string_chunks:
+            loader = Document(page_content=line)
+            documents_.append(loader)
+        embeddings = HuggingFaceEmbeddings()
+
+
+        db = FAISS.from_documents(documents_, embeddings)
+        print(type(db))
+        return db
+
+db = create_database()
 
 if "tokenizer" not in st.session_state:
     st.session_state["tokenizer"] = AutoTokenizer.from_pretrained(
@@ -14,7 +46,7 @@ if "tokenizer" not in st.session_state:
         "MBZUAI/LaMini-Flan-T5-783M"
     )
 
-st.title("B-Bot : Bibek's Personal Chatbot")
+st.title("BGPT : Bibek's Personal Chatbot")
 # Storing the chat
 if "generated" not in st.session_state:
     st.session_state["generated"] = []
@@ -30,25 +62,58 @@ def get_text():
 
 
 user_input = get_text()
+def wrap_text_preserve_newlines(text, width=110):
+    # Split the input text into lines based on newline characters
+    lines = text.split('\n')
+
+    # Wrap each line individually
+    wrapped_lines = [textwrap.fill(line, width=width) for line in lines]
+
+    # Join the wrapped lines back together using newline characters
+    wrapped_text = '\n'.join(wrapped_lines).replace("page_content=", "").replace("metadata={}", "")
+
+    return wrapped_text
 
 if user_input:
-    tokenizer = st.session_state["tokenizer"]
-    model = st.session_state["model"]
-    output = generate_response(user_input)
-    prompt_template = "\nBased on the above content, try to answer the following question.\n\n"
-    end_prompt = "Please make meaningful sentence and try to be descriptive as possible, ending with proper punctuations. If you think, there is good descriptive answers to the question from the above content, write sorry and advise them to contact Bibek directly.\n"  # NoQA"
-    short_response_template = "\nIf your response is very short like 1 or 2 sentence, add a followup sentence like 'Let me know if there's anything else I can help you with. or If there's anything else I can assist with, please don't hesitate to ask. I mean something similar in polite way."  # NoQA
-    input = output + prompt_template + user_input + end_prompt
+
+    tag = generate_tag(user_input)
+    
     start = time.time()
-    input_ids = tokenizer(
-        input,
-        return_tensors="pt",
-    ).input_ids
+    # print(tag)
+    if tag in ["greeting"]:
+        output = "Hello 👋! Thanks for visiting!\n I am BGPT! I am here to assist you in obtaining information about Bibek. Feel free to ask me any questions about Bibek. These are some sample questions:\n (I) Tell me about Bibek.\n (II) What skills does Bibek have?\n (III) What work experience does Bibek have?\n (IV) What is Bibek's educational background?\n (V) What awards has Bibek won?\n (VI) What projects have Bibek completed? &\n (VII) How can I contact Bibek?"
+    else:
+        tokenizer = st.session_state["tokenizer"]
+        model = st.session_state["model"]
+        docs = db.similarity_search(user_input)
+        output = wrap_text_preserve_newlines(str(docs[0]))
+        if tag in ["welcome", "thanks", "exit"]:
+            input = user_input    
+        elif tag in ["BibekBOT-introduction"]:
+            input = "I am BGPT, a large language model. I am here to assist you in obtaining information about Bibek. Feel free to ask me any questions about Bibek and I will make every effort to respond to all inquiries. These are some sample questions:\n (I) Tell me about Bibek.\n (II) What skills does Bibek have?\n (III) What work experience does Bibek have?\n (IV) What is Bibek's educational background?\n (V) What awards has Bibek won?\n (VI) What projects have Bibek completed? &\n (VII) How can I contact Bibek?. \n Can you paraphrase the above without changing the tone and contents." 
+        elif tag in ["decline"]:
+            input = "Okay, if there's anything else I can assist with, please don't hesitate to ask. \n Can you paraphrase the above without changing much content and tone."
+        else: 
+            # output = generate_response(user_input)
+            task_description_prompt = "I want you to act like my personal assistant chatbot named 'BGPT'. You are provided with some content and you will get one question. Try to answer the question in details based on the provided content. You may paraphrase the contents to reach your answer too. The below is the content: \n"
+            prompt_template = "\nBased on the above content, try to answer the following question.\n\n"
+            end_prompt = "\nPlease make meaningful sentence and try to be descriptive as possible responding with many sentences and ending with proper punctuations. If you think the content doesn't contain good answer to the question, give some polite respones telling them that you do not have specific response to the query and apologize and refer them to contact Bibek directly.\n"  # NoQA"
+            short_response_template = "\nIf your response is very short like 1 or 2 sentence, add a followup sentence like 'Let me know if there's anything else I can help you with. or If there's anything else I can assist with, please don't hesitate to ask. I mean something similar in polite way."  # NoQA
+
+            input = task_description_prompt + output + prompt_template + user_input + end_prompt
+            
+        input_ids = tokenizer(
+            input,
+            return_tensors="pt",
+        ).input_ids
+
+        outputs = model.generate(input_ids, max_length=512, do_sample=True)
+        output = tokenizer.decode(outputs[0]).strip("<pad></s>").strip()
 
-    outputs = model.generate(input_ids, max_length=512, do_sample=True)
-    output = tokenizer.decode(outputs[0]).strip('<pad></s>').strip()
     end = time.time()
 
+    # print(input)
+
     print("Time for model inference: ", end - start)
     # Checks for memory overflow
     if len(st.session_state.past) == 15: