Spaces:
Runtime error
Runtime error
File size: 5,893 Bytes
507d13c bcf7c58 507d13c bcf7c58 507d13c bcf7c58 99e964c 507d13c d118bdb 507d13c bcf7c58 507d13c 2c0084d 507d13c bcf7c58 d118bdb 507d13c d118bdb 507d13c 99e964c 2c0084d bcf7c58 507d13c bcf7c58 507d13c bcf7c58 507d13c 4fa708b d118bdb 4fa708b 2c0084d 4fa708b 2c0084d 4fa708b 2c0084d bcf7c58 2c0084d 4fa708b d118bdb 4fa708b 507d13c 4fa708b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | import cohere
import os
import pinecone
import uuid
from typing import List, Dict
from dotenv import load_dotenv
load_dotenv()
co = cohere.Client(os.environ["COHERE_API_KEY"])
pc = pinecone.Pinecone(api_key=os.environ["PINECONE_API_KEY"])
index = pc.Index("td-sec-embeddings")
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 = []
query_emb = co.embed(
texts=[query], model="embed-english-v3.0", input_type="search_query"
).embeddings
res = index.query(vector=query_emb, top_k=100, include_metadata=True)
docs_to_rerank = [match["metadata"] for match in res["matches"]]
rerank_results = co.rerank(
query=query,
documents=docs_to_rerank,
top_n=5,
model="rerank-english-v2.0",
)
docs_retrieved = []
for hit in rerank_results.results:
docs_retrieved.append(docs_to_rerank[hit.index])
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
self.docs = None
def send_initial_instructions(self):
response = self.co.chat_stream(
message="""You are an expert in TD Bank's annual reports and have access to the 2023 and 2022 annual report. Respond with a polite welcome message.""",
conversation_id=self.conversation_id,
)
return response
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,
conversation_id=self.conversation_id,
)
# If there are search queries, retrieve documents and respond
if response.search_queries:
print("Retrieving information...")
documents = self.retrieve_docs(response)
self.docs = {f"doc_{i}": document for i, document in enumerate(documents)}
response = self.co.chat_stream(
message=message,
documents=documents,
conversation_id=self.conversation_id,
)
for event in response:
yield event
# If there is no search query, directly respond
else:
response = self.co.chat_stream(
message=message,
conversation_id=self.conversation_id,
)
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))
return retrieved_docs
import gradio as gr
with gr.Blocks() as demo:
chatbot = gr.Chatbot()
msg = gr.Textbox()
clear = gr.Button("Clear")
cohere_chatbot_var = gr.State()
def user(user_message, history):
return "", history + [[user_message, None]]
def chat_function(history, cohere_chatbot):
if cohere_chatbot is None:
cohere_chatbot = Chatbot(co, index)
response = cohere_chatbot.send_initial_instructions()
history = [[None, ""]]
for event in response:
if event.event_type == "text-generation":
history[0][1] += str(event.text)
yield history, cohere_chatbot
return
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)
yield history, cohere_chatbot
# Citations
if event.event_type == "citation-generation":
if flag:
history[-1][1] += "\n\n**DOCUMENTS CONSULTED:**\n\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
# Make sure we run the thing once to initialize!
demo.load(
chat_function, [chatbot, cohere_chatbot_var], [chatbot, cohere_chatbot_var]
)
msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
chat_function, [chatbot, cohere_chatbot_var], [chatbot, cohere_chatbot_var]
)
clear.click(lambda: None, None, chatbot, queue=False)
demo.queue()
demo.launch()
|