Retrieval Augmented Generation (RAG)

1
# ! pip install cohere hnswlib unstructured
2
3
import cohere
4
5
co_chat = cohere.Client(
6
  api_key="AZURE_API_KEY_CHAT",
7
  base_url="AZURE_ENDPOINT_CHAT" # example: "https://cohere-command-r-plus-08-2024-xyz.eastus.models.ai.azure.com/"
8
)
9
10
co_embed = cohere.Client(
11
  api_key="AZURE_API_KEY_EMBED",
12
  base_url="AZURE_ENDPOINT_EMBED" # example: "https://cohere-embed-v3-multilingual-xyz.eastus.models.ai.azure.com/"
13
)
14
15
co_rerank = cohere.Client(
16
    api_key="AZURE_API_KEY_RERANK",
17
    base_url="AZURE_ENDPOINT_RERANK" # example: "https://cohere-rerank-v3-multilingual-xyz.eastus.models.ai.azure.com/"
18
)

1
documents = [
2
    {
3
        "title": "Tall penguins",
4
        "text": "Emperor penguins are the tallest."},
5
    {
6
        "title": "Penguin habitats",
7
        "text": "Emperor penguins only live in Antarctica."},
8
    {
9
        "title": "What are animals?",
10
        "text": "Animals are different from plants."}
11
]

1
message = "What are the tallest living penguins?"
2
3
response = co_chat.chat(
4
    message=message, 
5
    documents=documents
6
)
7
8
print("\nRESPONSE:\n")
9
print(response.text)
10
11
if response.citations:
12
    print("\nCITATIONS:\n")
13
    for citation in response.citations:
14
        print(citation)

1
import uuid
2
import yaml
3
import hnswlib
4
from typing import List, Dict
5
from unstructured.partition.html import partition_html
6
from unstructured.chunking.title import chunk_by_title

1
raw_documents = [
2
    {
3
        "title": "Crafting Effective Prompts",
4
        "url": "https://docs.cohere.com/docs/crafting-effective-prompts"},
5
    {
6
        "title": "Advanced Prompt Engineering Techniques",
7
        "url": "https://docs.cohere.com/docs/advanced-prompt-engineering-techniques"},
8
    {
9
        "title": "Prompt Truncation",
10
        "url": "https://docs.cohere.com/docs/prompt-truncation"},
11
    {
12
        "title": "Preambles",
13
        "url": "https://docs.cohere.com/docs/preambles"}
14
]

1
class Vectorstore:
2
3
    def __init__(self, raw_documents: List[Dict[str, str]]):
4
        self.raw_documents = raw_documents
5
        self.docs = []
6
        self.docs_embs = []
7
        self.retrieve_top_k = 10
8
        self.rerank_top_k = 3
9
        self.load_and_chunk()
10
        self.embed()
11
        self.index()
12
13
14
    def load_and_chunk(self) -> None:
15
        """
16
        Loads the text from the sources and chunks the HTML content.
17
        """
18
        print("Loading documents...")
19
20
        for raw_document in self.raw_documents:
21
            elements = partition_html(url=raw_document["url"])
22
            chunks = chunk_by_title(elements)
23
            for chunk in chunks:
24
                self.docs.append(
25
                    {
26
                        "title": raw_document["title"],
27
                        "text": str(chunk),
28
                        "url": raw_document["url"],
29
                    }
30
                )
31
32
    def embed(self) -> None:
33
        """
34
        Embeds the document chunks using the Cohere API.
35
        """
36
        print("Embedding document chunks...")
37
38
        batch_size = 90
39
        self.docs_len = len(self.docs)
40
        for i in range(0, self.docs_len, batch_size):
41
            batch = self.docs[i : min(i + batch_size, self.docs_len)]
42
            texts = [item["text"] for item in batch]
43
            docs_embs_batch = co_embed.embed(
44
                texts=texts,
45
                input_type="search_document"
46
            ).embeddings
47
            self.docs_embs.extend(docs_embs_batch)
48
49
    def index(self) -> None:
50
        """
51
        Indexes the document chunks for efficient retrieval.
52
        """
53
        print("Indexing document chunks...")
54
55
        self.idx = hnswlib.Index(space="ip", dim=1024)
56
        self.idx.init_index(max_elements=self.docs_len, ef_construction=512, M=64)
57
        self.idx.add_items(self.docs_embs, list(range(len(self.docs_embs))))
58
59
        print(f"Indexing complete with {self.idx.get_current_count()} document chunks.")
60
61
    def retrieve(self, query: str) -> List[Dict[str, str]]:
62
        """
63
        Retrieves document chunks based on the given query.
64
65
        Parameters:
66
        query (str): The query to retrieve document chunks for.
67
68
        Returns:
69
        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.
70
        """
71
72
        # Dense retrieval
73
        query_emb = co_embed.embed(
74
            texts=[query],
75
            input_type="search_query"
76
        ).embeddings
77
        
78
        doc_ids = self.idx.knn_query(query_emb, k=self.retrieve_top_k)[0][0]
79
80
        # Reranking
81
        docs_to_rerank = [self.docs[doc_id] for doc_id in doc_ids]
82
        yaml_docs = [yaml.dump(doc, sort_keys=False) for doc in docs_to_rerank] 
83
        rerank_results = co_rerank.rerank(
84
            query=query,
85
            documents=yaml_docs,
86
            top_n=self.rerank_top_k
87
        )
88
        doc_ids_reranked = [doc_ids[result.index] for result in rerank_results.results]
89
90
        docs_retrieved = []
91
        for doc_id in doc_ids_reranked:
92
            docs_retrieved.append(
93
                {
94
                    "title": self.docs[doc_id]["title"],
95
                    "text": self.docs[doc_id]["text"],
96
                    "url": self.docs[doc_id]["url"],
97
                }
98
            )
99
100
        return docs_retrieved

1
# Create an instance of the Vectorstore class with the given sources
2
vectorstore = Vectorstore(raw_documents)

1
vectorstore.retrieve("Prompting by giving examples")

1
def run_chatbot(message, chat_history=None):
2
        
3
    if chat_history is None:
4
        chat_history = []
5
    
6
    # Generate search queries, if any        
7
    response = co_chat.chat(
8
        message=message,
9
        search_queries_only=True,
10
        chat_history=chat_history,
11
    )
12
    
13
    search_queries = []
14
    for query in response.search_queries:
15
        search_queries.append(query.text)
16
17
    # If there are search queries, retrieve the documents
18
    if search_queries:
19
        print("Retrieving information...", end="")
20
21
        # Retrieve document chunks for each query
22
        documents = []
23
        for query in search_queries:
24
            documents.extend(vectorstore.retrieve(query))
25
26
        # Use document chunks to respond
27
        response = co_chat.chat(
28
            message=message,
29
            documents=documents,
30
            chat_history=chat_history
31
        )
32
33
    else:
34
        response = co_chat.chat(
35
            message=message,
36
            chat_history=chat_history
37
        )
38
        
39
    # Print the chatbot response, citations, and documents
40
    print("\nRESPONSE:\n")
41
    print(response.text)
42
        
43
    if response.citations:
44
        print("\nCITATIONS:\n")           
45
        for citation in response.citations:
46
            print(citation)
47
        print("\nDOCUMENTS:\n")           
48
        for document in response.documents:
49
            print(document)
50
            
51
    chat_history = response.chat_history
52
53
    return chat_history

1
chat_history = run_chatbot("Hello, I have a question")

1
chat_history = run_chatbot("What's the difference between zero-shot and few-shot prompting", chat_history)

1
chat_history = run_chatbot("What do you know about 5G networks?", chat_history)

1
print("Chat history:")
2
for c in chat_history:
3
    print(c, "\n")