Upload files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+chroma/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+.DS_Store
+__pycache__
+.venv

app.py

@@ -0,0 +1,30 @@
+import gradio as gr
+from src.ui import UI
+from src.chatbot import ChatBot
+
+def clear_session():
+    return "", []
+
+def add_query(chat_history, input):
+    if not input:
+        raise gr.Error("Please enter a question.")
+    chat_history.append((input, None))
+    return chat_history
+
+def response(chat_history, query):
+    res_msg, ref_docs = chatbot.generate_response(query, chat_history[:-1])
+    chat_history[-1] = (query, res_msg)
+    return "", chat_history, ref_docs
+
+if __name__ == "__main__":
+    demo, chatspace, ref_docs, text_input, clear_btn = UI.create_demo()
+
+    chatbot = ChatBot(is_debug=True)
+    with demo:
+        # Event handler for submitting text and generating response
+        text_input.submit(add_query, inputs=[chatspace, text_input], outputs=[chatspace], concurrency_limit=1).\
+            success(response, inputs=[chatspace, text_input], outputs=[text_input, chatspace, ref_docs])
+        clear_btn.click(clear_session, inputs=[], outputs=[text_input, chatspace])
+
+    demo.queue(api_open=False)
+    demo.launch()

chroma/c6cac3f8-bcae-47bc-a6c2-8fdc2770f19d/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2dbee596e8dcd6fe25e98be3ab98e6e07d4594cca8ff319396c9adabca1c7261
+size 37704000

chroma/c6cac3f8-bcae-47bc-a6c2-8fdc2770f19d/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:da7633500b48fa102f767e4c2b993bb473a37f64d13f12ccf076b4d058b671c9
+size 100

chroma/c6cac3f8-bcae-47bc-a6c2-8fdc2770f19d/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:383344cb187856235bccb00d648931a7988d828b7e52dedcb37e603489ae94ca
+size 346049

chroma/c6cac3f8-bcae-47bc-a6c2-8fdc2770f19d/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83c28735eeddff73cb1a12abd616afa354bf52e37ee424611241a641c30c913e
+size 24000

chroma/c6cac3f8-bcae-47bc-a6c2-8fdc2770f19d/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ef7ee997a8e34ddb4bc8487557d3299ebc0a8132ac19c9ef1b191a093eba3a69
+size 52152

chroma/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b30e09521687b19affc30b8f7082ceabd97a21b09735ecaf40610dc6cc1d230
+size 72232960

requirements.txt

@@ -0,0 +1,7 @@
+gradio==4.29.0
+langchain==0.1.17
+langchain_chroma==0.1.0
+langchain_community==0.0.36
+langchain_core==0.1.50
+langchain_openai==0.1.6
+langchain_text_splitters==0.0.1

src/chatbot.py

@@ -0,0 +1,128 @@
+import gradio as gr
+from typing import Any, Dict, List, Tuple
+from langchain_chroma import Chroma
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+from langchain.chains.history_aware_retriever import create_history_aware_retriever
+from langchain.chains.retrieval import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
+from langchain_core.documents import Document
+from langchain_core.retrievers import BaseRetriever
+import pandas as pd
+
+class CustomHandler(BaseCallbackHandler):
+    def __init__(self):
+        self.prompt = ""
+
+    def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> Any:
+        formatted_prompts = "\n".join(prompts)
+        self.prompt = formatted_prompts
+
+class CustomRetriever(BaseRetriever):
+    vectorstore: Chroma
+    comments: pd.DataFrame
+
+    def _get_relevant_documents(
+        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[Document]:   
+        docs = self.vectorstore.similarity_search(query)
+        matching_documents = []
+        for doc in docs:
+            post_id = int(doc.metadata['source'])
+            comment = self.comments.loc[self.comments['Post_ID'] == post_id, 'Comment_content'].values
+            query = doc.page_content.replace("Content: ", "User: ")
+            content = f"{query}\nAssistant: {comment[0]}"
+            matching_documents.append(
+                Document(
+                    page_content=content,
+                    metadata=doc.metadata
+                )
+            )
+
+        print(matching_documents)
+        return matching_documents
+
+class ChatBot:
+    def __init__(self, is_debug=False):
+        self.is_debug = is_debug
+        self.model = ChatOpenAI()
+        self.handler = CustomHandler()
+        self.embedding_function = OpenAIEmbeddings()
+        self.vectorstore = Chroma(
+            embedding_function=self.embedding_function,
+            collection_name="documents",
+            persist_directory="chroma", 
+        )
+        self.comments = pd.read_csv("data/comments.csv")
+        self.retriever = CustomRetriever(vectorstore=self.vectorstore, comments=self.comments)
+    
+    def create_chain(self):
+        qa_system_prompt = """
+        You are a helpful and joyous mental therapy assistant. Always answer as helpfully and cheerfully as possible, while being safe.
+        Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content.
+        Please ensure that your responses are socially unbiased and positive in nature.
+        If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. 
+        If you don't know the answer to a question, please don't share false information.
+
+        Here are a few examples of answers:   
+        {context}
+        
+        """
+        prompt = ChatPromptTemplate.from_messages([
+            ("system", qa_system_prompt),
+            MessagesPlaceholder(variable_name="chat_history"),
+            ("human", "{input}")
+        ])
+
+        chain = create_stuff_documents_chain(
+            llm=self.model,
+            prompt=prompt
+        )
+
+        retriever_prompt = ChatPromptTemplate.from_messages([
+            MessagesPlaceholder(variable_name="chat_history"),
+            ("human", "{input}"),
+            ("human", "Given the above conversation, generate a search query to look up in order to get information relevant to the conversation")
+        ])
+        history_aware_retriever = create_history_aware_retriever(
+            llm=self.model,
+            retriever=self.retriever,
+            prompt=retriever_prompt
+        )
+
+        retrieval_chain = create_retrieval_chain(
+            # retriever, Replace with History Aware Retriever
+            history_aware_retriever,
+            chain
+        )
+
+        return retrieval_chain
+
+    def process_chat_history(self, chat_history):
+        history = []
+        for (query, response) in chat_history:
+            history.append(HumanMessage(content=query))
+            history.append(AIMessage(content=response))
+        return history
+
+    def generate_response(self, query, chat_history):
+        if not input:
+            raise gr.Error("Please enter a question.")
+
+        history = self.process_chat_history(chat_history)
+        conversational_chain = self.create_chain()
+        response = conversational_chain.invoke(
+            {
+                "input": query,
+                "chat_history": history,
+            }, 
+            config={"callbacks": [self.handler]}
+        )["answer"]
+
+        references = self.handler.prompt if self.is_debug else "This is for debugging purposes only."
+        return response, references

src/ui.py

@@ -0,0 +1,28 @@
+import gradio as gr
+
+class UI:
+    @staticmethod
+    def feedback(data: gr.LikeData):
+        if data.liked:
+            print("You upvoted this response: " + data.value)
+        else:
+            print("You downvoted this response: " + data.value)
+
+    @staticmethod
+    def create_demo():
+        demo = gr.Blocks(title= "Chatbot", theme="Soft")
+        with demo:
+            with gr.Tab("Chat"):
+                chatbot = gr.Chatbot(value=[], elem_id='chatbot')
+                chatbot.like(UI.feedback, None, None)
+
+                text_input = gr.Textbox(
+                    show_label=False,
+                    placeholder="Ask me anything!",
+                container=False)
+
+                clear_btn = gr.Button("🧹 Clear")
+            with gr.Tab("Prompt"):
+                ref_docs = gr.Textbox(label='References', lines=25)
+
+        return demo, chatbot, ref_docs, text_input, clear_btn