add chatbot (#2)

* update requirements

* Add chatbot functionality

* isort

* put embedding caching in docparser, add error handling

* add typehints

* black

* package the project

* isort

buster/chatbot.py

@@ -0,0 +1,81 @@
+import logging
+import pickle
+
+import numpy as np
+import openai
+import pandas as pd
+from docparser import EMBEDDING_MODEL
+from openai.embeddings_utils import cosine_similarity, get_embedding
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# search through the reviews for a specific product
+def rank_documents(df: pd.DataFrame, query: str, top_k: int = 3) -> pd.DataFrame:
+    product_embedding = get_embedding(
+        query,
+        engine=EMBEDDING_MODEL,
+    )
+    df["similarity"] = df.embedding.apply(lambda x: cosine_similarity(x, product_embedding))
+
+    if top_k == -1:
+        # return all results
+        n = len(df)
+
+    results = df.sort_values("similarity", ascending=False).head(top_k)
+    return results
+
+
+def engineer_prompt(question: str, documents: list[str]) -> str:
+    return " ".join(documents) + "\nNow answer the following question:\n" + question
+
+
+def get_gpt_response(question: str, df) -> str:
+    # rank the documents, get the highest scoring doc and generate the prompt
+    candidates = rank_documents(df, query=question, top_k=1)
+    documents = candidates.documents.to_list()
+    prompt = engineer_prompt(question, documents)
+
+    logger.info(f"querying GPT...")
+    logger.info(f"User Question:\n{question}")
+    # Call the API to generate a response
+    try:
+        response = openai.Completion.create(
+            engine="text-davinci-003",
+            prompt=prompt,
+            max_tokens=200,
+            #  temperature=0,
+            #  top_p=0,
+            frequency_penalty=1,
+            presence_penalty=1,
+        )
+
+        # Get the response text
+        response_text = response["choices"][0]["text"]
+        logger.info(
+            f"""
+        GPT Response:\n{response_text}
+        """
+        )
+        return response_text
+    except Exception as e:
+        import traceback
+
+        logging.error(traceback.format_exc())
+        return "Oops, something went wrong. Try again later!"
+
+
+def load_embeddings(path: str) -> pd.DataFrame:
+    logger.info(f"loading embeddings from {path}...")
+    df = pd.read_csv(path)
+    df["embedding"] = df.embedding.apply(eval).apply(np.array)
+    logger.info(f"embeddings loaded.")
+    return df
+
+
+if __name__ == "__main__":
+    # we generate the embeddings using docparser.py
+    df = load_embeddings("data/document_embeddings.csv")
+
+    question = "Where should I put my datasets when I am running a job?"
+    response = get_gpt_response(question, df)

docparser.py → buster/docparser.py

@@ -2,7 +2,13 @@ import glob
 import os
 import pickle
 
+import pandas as pd
+import tiktoken
 from bs4 import BeautifulSoup
+from openai.embeddings_utils import cosine_similarity, get_embedding
+
+EMBEDDING_MODEL = "text-embedding-ada-002"
+EMBEDDING_ENCODING = "cl100k_base"  # this the encoding for text-embedding-ada-002
 
 
 def get_all_sections(root_dir: str, max_section_length: int = 3000) -> list[str]:
@@ -53,6 +59,35 @@ def read_sections(filepath: str) -> list[str]:
     return sections
 
 
+def load_documents(fname: str) -> pd.DataFrame:
+    df = pd.DataFrame()
+
+    with open(fname, "rb") as fp:
+        documents = pickle.load(fp)
+    df["documents"] = documents
+    return df
+
+
+def compute_n_tokens(df: pd.DataFrame) -> pd.DataFrame:
+    encoding = tiktoken.get_encoding(EMBEDDING_ENCODING)
+    df["n_tokens"] = df.documents.apply(lambda x: len(encoding.encode(x)))
+    return df
+
+
+def precompute_embeddings(df: pd.DataFrame) -> pd.DataFrame:
+    df["embedding"] = df.documents.apply(lambda x: get_embedding(x, engine=EMBEDDING_MODEL))
+    return df
+
+
+def generate_embeddings(filepath: str, output_csv: str) -> pd.DataFrame:
+    # Get all documents and precompute their embeddings
+    df = load_documents(filepath)
+    df = compute_n_tokens(df)
+    df = precompute_embeddings(df)
+    df.to_csv(output_csv)
+    return df
+
+
 if __name__ == "__main__":
     root_dir = "/home/hadrien/perso/mila-docs/output/"
     save_filepath = os.path.join(root_dir, "sections.pkl")
@@ -63,3 +98,6 @@ if __name__ == "__main__":
 
     # How to load
     sections = read_sections(save_filepath)
+
+    # precopmute the document embeddings
+    df = generate_embeddings(filepath=save_filepath, output_csv="data/document_embeddings.csv")

pyproject.toml

@@ -3,11 +3,11 @@ requires = ["setuptools", "setuptools-scm"]
 build-backend = "setuptools.build_meta"
 
 [project]
-name = "mila-cluster-chatbot"
+name = "buster"
 version = "0.0.1"
-description = "Chatbot to answer Mila cluster questions"
+description = "buster the bot for the mila cluster"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.9"
 dynamic = ["dependencies"]
 
 [tool.setuptools.dynamic]

requirements.txt

@@ -1,3 +1,4 @@
 pandas
 openai
 numpy
+tiktoken