Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ 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
+2023-annual-report-e.pdf filter=lfs diff=lfs merge=lfs -text
+ar2022-Complete-Report.pdf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,161 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+.python-version

2023-annual-report-e.pdf

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

README.md

@@ -1,12 +1,8 @@
 ---
-title: Secsplorer
-emoji: 🏃
-colorFrom: green
-colorTo: purple
+title: secsplorer
+app_file: simple_script.py
 sdk: gradio
 sdk_version: 4.13.0
-app_file: app.py
-pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# secsplorer
+A RAG ChatBot POC that works with a company's annual reports

ar2022-Complete-Report.pdf

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

docs/milestones.md

@@ -0,0 +1,18 @@
+# Milestone Plan
+
+Initial thoughts:
+
+The Cohere docs already explain how to create a RAG chatbot on your own machine. To make something awesome, it should come with a decent UI. Can be simple Gradio or Streamlit for now, but it should be there.
+
+Let's say the final goal is to have a bot that can answer questions based on TD's last 5 annual reports. One thing to note here is that makes the document situation static, whereas a company's true knowledge base will be dynamic. Anyway, for a POC it should be fine.
+
+-> Final goal is a website that displays a chatbot which you can ask questions about TD. How do we get there in a way that is fast and efficient?
+
+1. Get Cohere's _simple_ example running _locally_. Involves basic understanding, and having an API key.
+2. Put the simple example into a WebUI (Gradio?)
+3. Now rejig the simple example so it can work with PDFs
+4. Host it somewhere!
+
+Obviously, we could switch around 2 and 3. I'm not sure if either is better. Having the GUI first means we get to something sooner that's end-to-end. Doing the PDFs first means we tackle a slightly more difficult tech problem sooner. 🤷‍♂️.
+
+I'm leaving the hosting for last in case the prior steps inform strongly how that hosting should happen in the end.

gradio_tests.py

@@ -0,0 +1,11 @@
+import time
+import gradio as gr
+
+
+def slow_echo(message, history):
+    for i in range(len(message)):
+        time.sleep(0.3)
+        yield "You typed: " + message[: i + 1]
+
+
+gr.ChatInterface(slow_echo).launch()

modal_gradio_test.py

@@ -0,0 +1,26 @@
+from modal import Stub, Image, asgi_app
+
+from fastapi import FastAPI
+
+
+image = Image.debian_slim("3.11").pip_install(
+    "gradio",
+)
+
+stub = Stub("secsplorer", image=image)
+
+web_app = FastAPI()
+
+
+@stub.function()
+@asgi_app()
+def fastapi_app():
+    import gradio as gr
+    from gradio.routes import mount_gradio_app
+
+    def chat_function(message, history):
+        yield "Foo!"
+
+    interface = gr.ChatInterface(chat_function)
+
+    return mount_gradio_app(app=web_app, blocks=interface, path="/")

modal_script.py

@@ -0,0 +1,171 @@
+from modal import Stub, Image, Secret, asgi_app, method
+from urllib.request import urlretrieve
+from fastapi import FastAPI
+from typing import List, Dict
+
+
+image = Image.debian_slim("3.11").pip_install(
+    "cohere",
+    "gradio",
+    "pinecone-client",
+)
+
+stub = Stub("secsplorer", image=image)
+
+web_app = FastAPI()
+
+
+@stub.function(
+    secrets=[Secret.from_name("cohere-api-key"), Secret.from_name("pinecone-api-key")]
+)
+@asgi_app()
+def fastapi_app():
+    import cohere
+    import pinecone
+    import os
+    import uuid
+
+    import gradio as gr
+    from gradio.routes import mount_gradio_app
+
+    # print("Connecting to cohere client")
+    # co = cohere.Client(os.environ["COHERE_API_KEY"])
+    # print("Done")
+    # # pinecone.init(api_key=os.environ["PINECONE_API_KEY"], environment="gcp-starter")
+    # # index = pinecone.Index(index_name="td-sec-embeddings")
+    # index = None
+
+    # def retrieve(
+    #     index: pinecone.Index, query: str, co: cohere.Client
+    # ) -> List[Dict[str, str]]:
+    #     """
+    #     Retrieves documents based on the given query.
+
+    #     Parameters:
+    #     query (str): The query to retrieve documents for.
+
+    #     Returns:
+    #     List[Dict[str, str]]: A list of dictionaries representing the retrieved  documents, with 'title', 'snippet', and 'url' keys.
+    #     """
+    #     docs_retrieved = []
+
+    #     print(f"Calling retrieve for '{query}'")
+    #     print("Embedding the query")
+    #     query_emb = co.embed(
+    #         texts=[query], model="embed-english-v3.0", input_type="search_query"
+    #     ).embeddings
+
+    #     print("Querying pinecone")
+    #     res = index.query(query_emb, top_k=10, include_metadata=True)
+    #     print("Preparing to rerank")
+    #     docs_to_rerank = [match["metadata"] for match in res["matches"]]
+
+    #     rerank_results = co.rerank(
+    #         query=query,
+    #         documents=docs_to_rerank,
+    #         top_n=3,
+    #         model="rerank-english-v2.0",
+    #     )
+
+    #     docs_retrieved = []
+    #     for hit in rerank_results:
+    #         docs_retrieved.append(docs_to_rerank[hit.index])
+
+    #     print("Returning retrieved docs")
+    #     return docs_retrieved
+
+    # class Chatbot:
+    #     def __init__(self, co: cohere.Client, index: pinecone.Index):
+    #         self.index = index
+    #         self.conversation_id = str(uuid.uuid4())
+    #         self.co = co
+
+    #     def generate_response(self, message: str):
+    #         """
+    #         Generates a response to the user's message.
+
+    #         Parameters:
+    #         message (str): The user's message.
+
+    #         Yields:
+    #         Event: A response event generated by the chatbot.
+
+    #         Returns:
+    #         List[Dict[str, str]]: A list of dictionaries representing the retrieved documents.
+
+    #         """
+
+    #         # Generate search queries (if any)
+    #         response = self.co.chat(message=message, search_queries_only=True)
+
+    #         # If there are search queries, retrieve documents and respond
+    #         if response.search_queries:
+    #             print("Retrieving information")
+
+    #             documents = self.retrieve_docs(response)
+
+    #             response = self.co.chat(
+    #                 message=message,
+    #                 documents=documents,
+    #                 conversation_id=self.conversation_id,
+    #                 stream=True,
+    #             )
+    #             for event in response:
+    #                 yield event
+
+    #         # If there is no search query, directly respond
+    #         else:
+    #             response = self.co.chat(
+    #                 message=message, conversation_id=self.conversation_id, stream=True
+    #             )
+    #             for event in response:
+    #                 yield event
+
+    #     def retrieve_docs(self, response) -> List[Dict[str, str]]:
+    #         """
+    #         Retrieves documents based on the search queries in the response.
+
+    #         Parameters:
+    #         response: The response object containing search queries.
+
+    #         Returns:
+    #         List[Dict[str, str]]: A list of dictionaries representing the retrieved documents.
+
+    #         """
+    #         # Get the query(s)
+
+    #         queries = []
+    #         for search_query in response.search_queries:
+    #             queries.append(search_query["text"])
+
+    #         # Retrieve documents for each query
+    #         retrieved_docs = []
+    #         for query in queries:
+    #             retrieved_docs.extend(retrieve(self.index, query, self.co))
+
+    #         return retrieved_docs
+
+    # chatbot = Chatbot(co, index)
+
+    def chat_function(message, history):
+        return "Foo!"
+        # flag = False
+        # reply = ""
+        # for event in chatbot.generate_response(message):
+        #     if event.event_type == "text-generation":
+        #         reply += str(event.text)
+        #         yield reply
+
+        #     # Citations
+        #     if event.event_type == "citation-generation":
+        #         if not flag:
+        #             reply += "\n\nCITATIONS:\n\n"
+        #             yield reply
+        #             flag = True
+        #         reply += str(event.citations) + "\n"
+        #         yield reply
+
+    interface = gr.ChatInterface(chat_function)
+
+    print("All ready!")
+    return mount_gradio_app(app=web_app, blocks=interface, path="/")

requirements.txt

@@ -0,0 +1,101 @@
+aiofiles==23.2.1
+aiohttp==3.9.1
+aiosignal==1.3.1
+altair==5.2.0
+annotated-types==0.6.0
+anyio==4.2.0
+attrs==23.2.0
+backoff==2.2.1
+beautifulsoup4==4.12.2
+certifi==2023.11.17
+chardet==5.2.0
+charset-normalizer==3.3.2
+click==8.1.7
+cohere==4.40
+colorama==0.4.6
+contourpy==1.2.0
+cycler==0.12.1
+dataclasses-json==0.6.3
+emoji==2.9.0
+fastapi==0.108.0
+fastavro==1.9.2
+ffmpy==0.3.1
+filelock==3.13.1
+filetype==1.2.0
+fonttools==4.47.0
+frozenlist==1.4.1
+fsspec==2023.12.2
+gradio==4.13.0
+gradio_client==0.8.0
+h11==0.14.0
+hnswlib==0.8.0
+httpcore==1.0.2
+httpx==0.26.0
+huggingface-hub==0.20.2
+idna==3.6
+importlib-metadata==6.11.0
+importlib-resources==6.1.1
+Jinja2==3.1.2
+joblib==1.3.2
+jsonpath-python==1.0.6
+jsonschema==4.20.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+langdetect==1.0.9
+lxml==5.0.1
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+marshmallow==3.20.1
+matplotlib==3.8.2
+mdurl==0.1.2
+multidict==6.0.4
+mypy-extensions==1.0.0
+nltk==3.8.1
+numpy==1.26.3
+orjson==3.9.10
+packaging==23.2
+pandas==2.1.4
+pillow==10.2.0
+pydantic==2.5.3
+pydantic_core==2.14.6
+pydub==0.25.1
+Pygments==2.17.2
+PyMuPDF==1.23.8
+PyMuPDFb==1.23.7
+pyparsing==3.1.1
+python-dateutil==2.8.2
+python-dotenv==1.0.0
+python-iso639==2024.1.2
+python-magic==0.4.27
+python-multipart==0.0.6
+pytz==2023.3.post1
+PyYAML==6.0.1
+rapidfuzz==3.6.1
+referencing==0.32.1
+regex==2023.12.25
+requests==2.31.0
+rich==13.7.0
+rpds-py==0.16.2
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.0
+soupsieve==2.5
+starlette==0.32.0.post1
+tabulate==0.9.0
+tomlkit==0.12.0
+toolz==0.12.0
+tqdm==4.66.1
+typer==0.9.0
+typing-inspect==0.9.0
+typing_extensions==4.9.0
+tzdata==2023.4
+unstructured==0.11.8
+unstructured-client==0.15.1
+urllib3==2.1.0
+uvicorn==0.25.0
+websockets==11.0.3
+wrapt==1.16.0
+yarl==1.9.4
+zipp==3.17.0
+modal

secsplorer/__init__.py

@@ -0,0 +1,4 @@
+from .documents import Documents
+from .chatbot import Chatbot
+
+__all__ = [Chatbot, Documents]

secsplorer/chatbot.py

@@ -0,0 +1,75 @@
+import uuid
+import cohere
+from secsplorer import Documents
+from typing import List, Dict
+
+
+class Chatbot:
+    def __init__(self, co: cohere.Client, docs: Documents):
+        self.docs = docs
+        self.conversation_id = str(uuid.uuid4())
+        self.co = co
+
+    def generate_response(self, message: str):
+        """
+        Generates a response to the user's message.
+
+        Parameters:
+        message (str): The user's message.
+
+        Yields:
+        Event: A response event generated by the chatbot.
+
+        Returns:
+        List[Dict[str, str]]: A list of dictionaries representing the retrieved documents.
+
+        """
+
+        # Generate search queries (if any)
+        response = self.co.chat(message=message, search_queries_only=True)
+
+        # If there are search queries, retrieve documents and respond
+        if response.search_queries:
+            print("Retrieving information...")
+
+            documents = self.retrieve_docs(response)
+
+            response = self.co.chat(
+                message=message,
+                documents=documents,
+                conversation_id=self.conversation_id,
+                stream=True,
+            )
+            for event in response:
+                yield event
+
+        # If there is no search query, directly respond
+        else:
+            response = self.co.chat(
+                message=message, conversation_id=self.conversation_id, stream=True
+            )
+            for event in response:
+                yield event
+
+    def retrieve_docs(self, response) -> List[Dict[str, str]]:
+        """
+        Retrieves documents based on the search queries in the response.
+
+        Parameters:
+        response: The response object containing search queries.
+
+        Returns:
+        List[Dict[str, str]]: A list of dictionaries representing the retrieved documents.
+
+        """
+        # Get the query(s)
+        queries = []
+        for search_query in response.search_queries:
+            queries.append(search_query["text"])
+
+        # Retrieve documents for each query
+        retrieved_docs = []
+        for query in queries:
+            retrieved_docs.extend(self.docs.retrieve(query, self.co))
+
+        return retrieved_docs

secsplorer/documents.py

@@ -0,0 +1,110 @@
+from typing import List, Dict
+from unstructured.partition.pdf import partition_pdf
+from unstructured.chunking.title import chunk_by_title
+import hnswlib
+import cohere
+
+
+class Documents:
+    def __init__(self, co: cohere.Client, sources: List[Dict[str, str]]):
+        self.co = co
+        self.sources = sources
+        self.docs = []
+        self.docs_embs = []
+        self.retrieve_top_k = 10
+        self.rerank_top_k = 3
+        self.load()
+        self.embed(co)
+        self.index()
+
+    def load(self) -> None:
+        """
+        Loads the documents from the sources and chunks the HTML content.
+        """
+        print("Loading documents...")
+
+        for source in self.sources:
+            # elements = partition_html(url=source["url"])
+            # elements = partition_pdf(url=source["url"])
+            elements = partition_pdf(filename=source["filename"])
+            chunks = chunk_by_title(elements)
+            for chunk in chunks:
+                self.docs.append(
+                    {
+                        "title": source["title"],
+                        "text": str(chunk),
+                        "url": source["url"],
+                    }
+                )
+
+    def embed(self) -> None:
+        """
+        Embeds the documents using the Cohere API.
+        """
+        print("Embedding documents...")
+
+        batch_size = 90
+        self.docs_len = len(self.docs)
+
+        for i in range(0, self.docs_len, batch_size):
+            batch = self.docs[i : min(i + batch_size, self.docs_len)]
+            texts = [item["text"] for item in batch]
+            docs_embs_batch = self.co.embed(
+                texts=texts, model="embed-english-v3.0", input_type="search_document"
+            ).embeddings
+            self.docs_embs.extend(docs_embs_batch)
+
+    def index(self) -> None:
+        """
+        Indexes the documents for efficient retrieval.
+        """
+        print("Indexing documents...")
+
+        self.index = hnswlib.Index(space="ip", dim=1024)
+        self.index.init_index(max_elements=self.docs_len, ef_construction=512, M=64)
+        self.index.add_items(self.docs_embs, list(range(len(self.docs_embs))))
+
+        print(f"Indexing complete with {self.index.get_current_count()} documents.")
+
+    def retrieve(self, query: str) -> List[Dict[str, str]]:
+        """
+        Retrieves documents based on the given query.
+
+        Parameters:
+        query (str): The query to retrieve documents for.
+
+        Returns:
+        List[Dict[str, str]]: A list of dictionaries representing the retrieved  documents, with 'title', 'snippet', and 'url' keys.
+        """
+        docs_retrieved = []
+        query_emb = self.co.embed(
+            texts=[query], model="embed-english-v3.0", input_type="search_query"
+        ).embeddings
+
+        doc_ids = self.index.knn_query(query_emb, k=self.retrieve_top_k)[0][0]
+
+        docs_to_rerank = []
+        for doc_id in doc_ids:
+            docs_to_rerank.append(self.docs[doc_id]["text"])
+
+        rerank_results = self.co.rerank(
+            query=query,
+            documents=docs_to_rerank,
+            top_n=self.rerank_top_k,
+            model="rerank-english-v2.0",
+        )
+
+        doc_ids_reranked = []
+        for result in rerank_results:
+            doc_ids_reranked.append(doc_ids[result.index])
+
+        for doc_id in doc_ids_reranked:
+            docs_retrieved.append(
+                {
+                    "title": self.docs[doc_id]["title"],
+                    "text": self.docs[doc_id]["text"],
+                    "url": self.docs[doc_id]["url"],
+                }
+            )
+
+        return docs_retrieved

setup.py

@@ -0,0 +1,10 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="secsplorer",
+    version="0.1",
+    packages=find_packages(),
+    author="Clemens Adolphs",
+    author_email="clemens.adolphs@gmail.com",
+    description="The private python package for the secsplorer",
+)

simple_script.py

@@ -0,0 +1,238 @@
+import cohere
+import os
+import pinecone
+import uuid
+
+from typing import List, Dict
+
+# from unstructured.chunking.title import chunk_by_title
+# from unstructured.partition.pdf import partition_pdf
+from dotenv import load_dotenv
+
+
+load_dotenv()
+
+co = cohere.Client(os.environ["COHERE_API_KEY"])
+
+pinecone.init(api_key=os.environ["PINECONE_API_KEY"], environment="gcp-starter")
+
+index = pinecone.Index("td-sec-embeddings")
+
+from typing import List, Dict
+
+# from unstructured.partition.pdf import partition_pdf
+# from unstructured.chunking.title import chunk_by_title
+
+import cohere
+
+sources = [
+    {
+        "title": "2023",
+        "url": "https://www.td.com/content/dam/tdcom/canada/about-td/pdf/quarterly-results/2023/2023-annual-report-e.pdf",
+        "filename": "/Users/clemensadolphs/git-personal/secsplorer/2023-annual-report-e.pdf",
+    },
+    # {
+    #     "title": "2022",
+    #     "url": "https://www.td.com/document/PDF/ar2022/ar2022-Complete-Report.pdf",
+    #     "filename": "/Users/clemensadolphs/git-personal/secsplorer/2023-annual-report-e.pdf",
+    # },
+]
+
+
+def load() -> List[Dict[str, str]]:
+    """
+    Loads the documents from the sources and chunks the HTML content.
+    """
+    print("Loading documents...")
+    docs = []
+    for source in sources:
+        elements = partition_pdf(filename=source["filename"])
+        chunks = chunk_by_title(elements)
+        for chunk in chunks:
+            docs.append(
+                {
+                    "title": source["title"],
+                    "text": str(chunk),
+                    "url": source["url"],
+                }
+            )
+    return docs
+
+
+def embed(docs: List[Dict[str, str]]) -> List[List[float]]:
+    """
+    Embeds the documents using the Cohere API.
+    """
+    print("Embedding documents...")
+
+    batch_size = 90
+    docs_len = len(docs)
+    docs_embs = []
+
+    for i in range(0, docs_len, batch_size):
+        batch = docs[i : min(i + batch_size, docs_len)]
+        texts = [item["text"] for item in batch]
+        docs_embs_batch = co.embed(
+            texts=texts, model="embed-english-v3.0", input_type="search_document"
+        ).embeddings
+        docs_embs.extend(docs_embs_batch)
+    return docs_embs
+
+
+def update_index(
+    index: pinecone.Index, docs: List[Dict[str, str]], docs_embs: List[List[float]]
+) -> None:
+    """
+    Indexes the documents for efficient retrieval.
+    """
+    batch_size = 100
+
+    ids = [str(i) for i in range(len(docs))]
+
+    to_upsert = list(zip(ids, docs_embs, docs))
+
+    for i in range(0, len(docs), batch_size):
+        i_end = min(i + batch_size, len(docs))
+        index.upsert(vectors=to_upsert[i:i_end])
+
+
+def retrieve(index: pinecone.Index, query: str) -> List[Dict[str, str]]:
+    """
+    Retrieves documents based on the given query.
+
+    Parameters:
+    query (str): The query to retrieve documents for.
+
+    Returns:
+    List[Dict[str, str]]: A list of dictionaries representing the retrieved  documents, with 'title', 'snippet', and 'url' keys.
+    """
+    docs_retrieved = []
+
+    print(f"Calling retrieve for '{query}'")
+    print("Embedding the query")
+    query_emb = co.embed(
+        texts=[query], model="embed-english-v3.0", input_type="search_query"
+    ).embeddings
+
+    print("Querying pinecone")
+    res = index.query(query_emb, top_k=10, include_metadata=True)
+    print("Preparing to rerank")
+    docs_to_rerank = [match["metadata"] for match in res["matches"]]
+
+    rerank_results = co.rerank(
+        query=query,
+        documents=docs_to_rerank,
+        top_n=3,
+        model="rerank-english-v2.0",
+    )
+
+    docs_retrieved = []
+    for hit in rerank_results:
+        docs_retrieved.append(docs_to_rerank[hit.index])
+
+    print("Returning retrieved docs")
+    return docs_retrieved
+
+
+# docs = load()
+# docs_embeds = embed(docs)
+# update_index(index, docs=docs, docs_embs=docs_embeds)
+
+
+class Chatbot:
+    def __init__(self, co: cohere.Client, index: pinecone.Index):
+        self.index = index
+        self.conversation_id = str(uuid.uuid4())
+        self.co = co
+
+    def generate_response(self, message: str):
+        """
+        Generates a response to the user's message.
+
+        Parameters:
+        message (str): The user's message.
+
+        Yields:
+        Event: A response event generated by the chatbot.
+
+        Returns:
+        List[Dict[str, str]]: A list of dictionaries representing the retrieved documents.
+
+        """
+
+        # Generate search queries (if any)
+        response = self.co.chat(message=message, search_queries_only=True)
+
+        # If there are search queries, retrieve documents and respond
+        if response.search_queries:
+            print("Retrieving information...")
+
+            documents = self.retrieve_docs(response)
+            print(f"Generating response with documents {documents}")
+            response = self.co.chat(
+                message=message,
+                documents=documents,
+                conversation_id=self.conversation_id,
+                stream=True,
+            )
+            for event in response:
+                yield event
+
+        # If there is no search query, directly respond
+        else:
+            response = self.co.chat(
+                message=message, conversation_id=self.conversation_id, stream=True
+            )
+            for event in response:
+                yield event
+
+    def retrieve_docs(self, response) -> List[Dict[str, str]]:
+        """
+        Retrieves documents based on the search queries in the response.
+
+        Parameters:
+        response: The response object containing search queries.
+
+        Returns:
+        List[Dict[str, str]]: A list of dictionaries representing the retrieved documents.
+
+        """
+        # Get the query(s)
+        print("Calling retrieve_docs")
+        queries = []
+        for search_query in response.search_queries:
+            queries.append(search_query["text"])
+        print(queries)
+
+        # Retrieve documents for each query
+        retrieved_docs = []
+        for query in queries:
+            retrieved_docs.extend(retrieve(self.index, query))
+
+        return retrieved_docs
+
+
+import gradio as gr
+
+chatbot = Chatbot(co, index)
+
+
+def chat_function(message, history):
+    flag = False
+    reply = ""
+    for event in chatbot.generate_response(message):
+        if event.event_type == "text-generation":
+            reply += str(event.text)
+            yield reply
+
+        # Citations
+        if event.event_type == "citation-generation":
+            if not flag:
+                reply += "\n\nCITATIONS:\n\n"
+                yield reply
+                flag = True
+            reply += str(event.citations) + "\n"
+            yield reply
+
+
+gr.ChatInterface(chat_function).launch(share=True)