Retrieval augmented generation (RAG) - Cohere on Azure AI Foundry

1
# %pip install cohere hnswlib unstructured
2
3
import cohere
4
5
co_chat = cohere.ClientV2(
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.ClientV2(
11
    api_key="AZURE_API_KEY_EMBED",
12
    base_url="AZURE_ENDPOINT_EMBED",  # example: "https://embed-v-4-0-xyz.eastus.models.ai.azure.com/"
13
)
14
15
co_rerank = cohere.ClientV2(
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
    {
7
        "title": "Penguin habitats",
8
        "text": "Emperor penguins only live in Antarctica.",
9
    },
10
    {
11
        "title": "What are animals?",
12
        "text": "Animals are different from plants.",
13
    },
14
]

1
message = "What are the tallest living penguins?"
2
3
response = co_chat.chat(
4
    model="model",  # Pass a dummy string
5
    messages=[{"role": "user", "content": message}],
6
    documents=[{"data": doc} for doc in documents],
7
)
8
9
print("\nRESPONSE:\n")
10
print(response.message.content[0].text)
11
12
if response.message.citations:
13
    print("\nCITATIONS:\n")
14
    for citation in response.message.citations:
15
        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
    {
7
        "title": "Advanced Prompt Engineering Techniques",
8
        "url": "https://docs.cohere.com/docs/advanced-prompt-engineering-techniques",
9
    },
10
    {
11
        "title": "Prompt Truncation",
12
        "url": "https://docs.cohere.com/docs/prompt-truncation",
13
    },
14
    {
15
        "title": "Preambles",
16
        "url": "https://docs.cohere.com/docs/preambles",
17
    },
18
]

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
    def load_and_chunk(self) -> None:
14
        """
15
        Loads the text from the sources and chunks the HTML content.
16
        """
17
        print("Loading documents...")
18
19
        for raw_document in self.raw_documents:
20
            elements = partition_html(url=raw_document["url"])
21
            chunks = chunk_by_title(elements)
22
            for chunk in chunks:
23
                self.docs.append(
24
                    {
25
                        "data": {
26
                            "title": raw_document["title"],
27
                            "text": str(chunk),
28
                            "url": raw_document["url"],
29
                        }
30
                    }
31
                )
32
33
    def embed(self) -> None:
34
        """
35
        Embeds the document chunks using the Cohere API.
36
        """
37
        print("Embedding document chunks...")
38
39
        batch_size = 90
40
        self.docs_len = len(self.docs)
41
        for i in range(0, self.docs_len, batch_size):
42
            batch = self.docs[i : min(i + batch_size, self.docs_len)]
43
            texts = [item["data"]["text"] for item in batch]
44
            docs_embs_batch = co_embed.embed(
45
                texts=texts,
46
                model="embed-v4.0",
47
                input_type="search_document",
48
                embedding_types=["float"],
49
            ).embeddings.float
50
            self.docs_embs.extend(docs_embs_batch)
51
52
    def index(self) -> None:
53
        """
54
        Indexes the document chunks for efficient retrieval.
55
        """
56
        print("Indexing document chunks...")
57
58
        self.idx = hnswlib.Index(space="ip", dim=1024)
59
        self.idx.init_index(
60
            max_elements=self.docs_len, ef_construction=512, M=64
61
        )
62
        self.idx.add_items(
63
            self.docs_embs, list(range(len(self.docs_embs)))
64
        )
65
66
        print(
67
            f"Indexing complete with {self.idx.get_current_count()} document chunks."
68
        )
69
70
    def retrieve(self, query: str) -> List[Dict[str, str]]:
71
        """
72
        Retrieves document chunks based on the given query.
73
74
        Parameters:
75
        query (str): The query to retrieve document chunks for.
76
77
        Returns:
78
        List[Dict[str, str]]: A list of dictionaries representing the retrieved document chunks, with 'title', 'text', and 'url' keys.
79
        """
80
81
        # Dense retrieval
82
        query_emb = co_embed.embed(
83
            texts=[query],
84
            model="embed-v4.0",
85
            input_type="search_query",
86
            embedding_types=["float"],
87
        ).embeddings.float
88
89
        doc_ids = self.idx.knn_query(
90
            query_emb, k=self.retrieve_top_k
91
        )[0][0]
92
93
        # Reranking
94
        docs_to_rerank = [
95
            self.docs[doc_id]["data"] for doc_id in doc_ids
96
        ]
97
        yaml_docs = [
98
            yaml.dump(doc, sort_keys=False) for doc in docs_to_rerank
99
        ]
100
        rerank_results = co_rerank.rerank(
101
            query=query,
102
            documents=yaml_docs,
103
            model="model",  # Pass a dummy string
104
            top_n=self.rerank_top_k,
105
        )
106
107
        doc_ids_reranked = [
108
            doc_ids[result.index] for result in rerank_results.results
109
        ]
110
111
        docs_retrieved = []
112
        for doc_id in doc_ids_reranked:
113
            docs_retrieved.append(self.docs[doc_id]["data"])
114
115
        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(query, messages=None):
2
    if messages is None:
3
        messages = []
4
5
    messages.append({"role": "user", "content": query})
6
7
    # Retrieve document chunks and format
8
    documents = vectorstore.retrieve(query)
9
    documents_formatted = []
10
    for doc in documents:
11
        documents_formatted.append({"data": doc})
12
13
    # Use document chunks to respond
14
    response = co_chat.chat(
15
        model="model",  # Pass a dummy string
16
        messages=messages,
17
        documents=documents_formatted,
18
    )
19
20
    # Print the chatbot response, citations, and documents
21
    print("\nRESPONSE:\n")
22
    print(response.message.content[0].text)
23
24
    if response.message.citations:
25
        print("\nCITATIONS:\n")
26
        for citation in response.message.citations:
27
            print("-" * 20)
28
            print(
29
                "start:",
30
                citation.start,
31
                "end:",
32
                citation.end,
33
                "text:",
34
                citation.text,
35
            )
36
            print("SOURCES:")
37
            print(citation.sources)
38
39
    # Add assistant response to messages
40
    messages.append(
41
        {
42
            "role": "assistant",
43
            "content": response.message.content[0].text,
44
        }
45
    )
46
47
    return messages

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

1
messages = run_chatbot("How to provide examples in prompts", messages)

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

1
for message in messages:
2
    print(message, "\n")