RAG With Chat Embed and Rerank via Pinecone

1
! pip install cohere hnswlib unstructured python-dotenv -q

1
import cohere
2
from pinecone import Pinecone, PodSpec
3
import uuid
4
import hnswlib
5
from typing import List, Dict
6
from unstructured.partition.html import partition_html
7
from unstructured.chunking.title import chunk_by_title
8
9
co = cohere.Client("COHERE_API_KEY") # Get your API key here: https://dashboard.cohere.com/api-keys
10
pc = Pinecone(api_key="PINECONE_API_KEY") # (get API key at app.pinecone.io)

1
import cohere
2
import os
3
import dotenv
4
5
dotenv.load_dotenv()
6
co = cohere.Client(os.getenv("COHERE_API_KEY"))
7
pc = Pinecone(
8
    api_key=os.getenv("PINECONE_API_KEY")
9
)

1
raw_documents = [
2
    {
3
        "title": "Text Embeddings",
4
        "url": "https://docs.cohere.com/docs/text-embeddings"},
5
    {
6
        "title": "Similarity Between Words and Sentences",
7
        "url": "https://docs.cohere.com/docs/similarity-between-words-and-sentences"},
8
    {
9
        "title": "The Attention Mechanism",
10
        "url": "https://docs.cohere.com/docs/the-attention-mechanism"},
11
    {
12
        "title": "Transformer Models",
13
        "url": "https://docs.cohere.com/docs/transformer-models"}
14
]

1
class Vectorstore:
2
    """
3
    A class representing a collection of documents indexed into a vectorstore.
4
5
    Parameters:
6
    raw_documents (list): A list of dictionaries representing the sources of the raw documents. Each dictionary should have 'title' and 'url' keys.
7
8
    Attributes:
9
    raw_documents (list): A list of dictionaries representing the raw documents.
10
    docs (list): A list of dictionaries representing the chunked documents, with 'title', 'text', and 'url' keys.
11
    docs_embs (list): A list of the associated embeddings for the document chunks.
12
    docs_len (int): The number of document chunks in the collection.
13
    idx (hnswlib.Index): The index used for document retrieval.
14
15
    Methods:
16
    load_and_chunk(): Loads the data from the sources and partitions the HTML content into chunks.
17
    embed(): Embeds the document chunks using the Cohere API.
18
    index(): Indexes the document chunks for efficient retrieval.
19
    retrieve(): Retrieves document chunks based on the given query.
20
    """
21
22
    def __init__(self, raw_documents: List[Dict[str, str]]):
23
        self.raw_documents = raw_documents
24
        self.docs = []
25
        self.docs_embs = []
26
        self.retrieve_top_k = 10
27
        self.rerank_top_k = 3
28
        self.load_and_chunk()
29
        self.embed()
30
        self.index()
31
32
33
    def load_and_chunk(self) -> None:
34
        """
35
        Loads the text from the sources and chunks the HTML content.
36
        """
37
        print("Loading documents...")
38
39
        for raw_document in self.raw_documents:
40
            elements = partition_html(url=raw_document["url"])
41
            chunks = chunk_by_title(elements)
42
            for chunk in chunks:
43
                self.docs.append(
44
                    {
45
                        "title": raw_document["title"],
46
                        "text": str(chunk),
47
                        "url": raw_document["url"],
48
                    }
49
                )
50
51
    def embed(self) -> None:
52
        """
53
        Embeds the document chunks using the Cohere API.
54
        """
55
        print("Embedding document chunks...")
56
57
        batch_size = 90
58
        self.docs_len = len(self.docs)
59
        for i in range(0, self.docs_len, batch_size):
60
            batch = self.docs[i : min(i + batch_size, self.docs_len)]
61
            texts = [item["text"] for item in batch]
62
            docs_embs_batch = co.embed(
63
                texts=texts, model="embed-english-v3.0", input_type="search_document"
64
            ).embeddings
65
            self.docs_embs.extend(docs_embs_batch)
66
67
    def index(self) -> None:
68
        """
69
        Indexes the documents for efficient retrieval.
70
        """
71
        print("Indexing documents...")
72
73
        index_name = 'rag-01'
74
75
        # If the index does not exist, we create it
76
        if index_name not in pc.list_indexes().names():
77
            pc.create_index(
78
                name=index_name,
79
                dimension=len(self.docs_embs[0]),
80
                metric="cosine",
81
                spec=PodSpec(
82
                    environment="gcp-starter"
83
                )
84
                )
85
86
        # connect to index
87
        self.idx = pc.Index(index_name)
88
89
        batch_size = 128
90
91
        ids = [str(i) for i in range(len(self.docs))]
92
        # create list of metadata dictionaries
93
        meta = self.docs
94
95
        # create list of (id, vector, metadata) tuples to be upserted
96
        to_upsert = list(zip(ids, self.docs_embs, meta))
97
98
        for i in range(0, len(self.docs), batch_size):
99
            i_end = min(i+batch_size, len(self.docs))
100
            self.idx.upsert(vectors=to_upsert[i:i_end])
101
102
        # let's view the index statistics
103
        print("Indexing complete")
104
105
106
    def retrieve(self, query: str) -> List[Dict[str, str]]:
107
        """
108
        Retrieves document chunks based on the given query.
109
110
        Parameters:
111
        query (str): The query to retrieve document chunks for.
112
113
        Returns:
114
        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.
115
        """
116
117
        docs_retrieved = []
118
        query_emb = co.embed(
119
            texts=[query], model="embed-english-v3.0", input_type="search_query"
120
        ).embeddings
121
122
123
        res = self.idx.query(vector=query_emb, top_k=self.retrieve_top_k, include_metadata=True)
124
        docs_to_rerank = [match['metadata']['text'] for match in res['matches']]
125
126
        rerank_results = co.rerank(
127
            query=query,
128
            documents=docs_to_rerank,
129
            top_n=self.rerank_top_k,
130
            model="rerank-english-v2.0",
131
        )
132
133
        docs_reranked = [res['matches'][result.index] for result in rerank_results.results]
134
135
        for doc in docs_reranked:
136
            docs_retrieved.append(doc['metadata'])
137
138
        return docs_retrieved

1
vectorstore = Vectorstore(raw_documents)

1
vectorstore.retrieve("multi-head attention definition")

1
class Chatbot:
2
    def __init__(self, vectorstore: Vectorstore):
3
        """
4
        Initializes an instance of the Chatbot class.
5
6
        Parameters:
7
        vectorstore (Vectorstore): An instance of the Vectorstore class.
8
9
        """
10
        self.vectorstore = vectorstore
11
        self.conversation_id = str(uuid.uuid4())
12
13
    def run(self):
14
        """
15
        Runs the chatbot application.
16
17
        """
18
        while True:
19
            # Get the user message
20
            message = input("User: ")
21
22
            # Typing "quit" ends the conversation
23
            if message.lower() == "quit":
24
              print("Ending chat.")
25
              break
26
            # else:                       # Uncomment for Google Colab to avoid printing the same thing twice
27
              # print(f"User: {message}") # Uncomment for Google Colab to avoid printing the same thing twice
28
29
            # Generate search queries (if any)
30
            response = co.chat(message=message,
31
                               model="command-r",
32
                               search_queries_only=True)
33
34
            # If there are search queries, retrieve document chunks and respond
35
            if response.search_queries:
36
                print("Retrieving information...", end="")
37
38
                # Retrieve document chunks for each query
39
                documents = []
40
                for query in response.search_queries:
41
                    documents.extend(self.vectorstore.retrieve(query.text))
42
43
                # Use document chunks to respond
44
                response = co.chat_stream(
45
                    message=message,
46
                    model="command-r",
47
                    documents=documents,
48
                    conversation_id=self.conversation_id,
49
                )
50
51
            # If there is no search query, directly respond
52
            else:
53
                response = co.chat_stream(
54
                    message=message,
55
                    model="command-r",
56
                    conversation_id=self.conversation_id,
57
                )
58
59
            # Print the chatbot response, citations, and documents
60
            print("\nChatbot:")
61
            citations = []
62
            cited_documents = []
63
64
            # Display response
65
            for event in response:
66
                if event.event_type == "text-generation":
67
                    print(event.text, end="")
68
                elif event.event_type == "citation-generation":
69
                    citations.extend(event.citations)
70
                elif event.event_type == "search-results":
71
                    cited_documents = event.documents
72
73
            # Display citations and source documents
74
            if citations:
75
              print("\n\nCITATIONS:")
76
              for citation in citations:
77
                print(citation)
78
79
              print("\nDOCUMENTS:")
80
              for document in cited_documents:
81
                print(document)
82
83
            print(f"\n{'-'*100}\n")

1
chatbot = Chatbot(vectorstore)
2
3
chatbot.run()