Upload folder using huggingface_hub

simple_script.py

@@ -55,6 +55,7 @@ class Chatbot:
         self.index = index
         self.conversation_id = str(uuid.uuid4())
         self.co = co
+        self.docs = None
 
     def generate_response(self, message: str):
         """
@@ -80,6 +81,8 @@ class Chatbot:
 
             documents = self.retrieve_docs(response)
 
+            self.docs = {f"doc_{i}": document for i, document in enumerate(documents)}
+
             response = self.co.chat(
                 message=message,
                 documents=documents,
@@ -134,13 +137,15 @@ with gr.Blocks() as demo:
         return "", history + [[user_message, None]]
 
     def chat_function(history, cohere_chatbot):
-        flag = False
         if cohere_chatbot is None:
             cohere_chatbot = Chatbot(co, index)
 
         message = history[-1][0]
         history[-1][1] = ""
 
+        documents_used = set()
+        flag = True
+
         for event in cohere_chatbot.generate_response(message):
             if event.event_type == "text-generation":
                 history[-1][1] += str(event.text)
@@ -148,12 +153,19 @@ with gr.Blocks() as demo:
 
             # Citations
             if event.event_type == "citation-generation":
-                if not flag:
-                    history[-1][1] += "\n\nCITATIONS:\n\n"
+                if flag:
+                    history[-1][1] += "\n\n**DOCUMENTS CONSULTED:**\n\n"
                     yield history, cohere_chatbot
-                    flag = True
-                history[-1][1] += str(event.citations) + "\n"
-                yield history, cohere_chatbot
+                    flag = False
+
+                for citation in event.citations:
+                    documents_used.update(citation["document_ids"])
+
+        urls_used = set(cohere_chatbot.docs[doc_id]["url"] for doc_id in documents_used)
+
+        for url in sorted(urls_used):
+            history[-1][1] += f"* {url}\n"
+            yield history, cohere_chatbot
 
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         chat_function, [chatbot, cohere_chatbot_var], [chatbot, cohere_chatbot_var]