Creating a QA Bot From Technical Documentation

1
%%capture
2
!pip install cohere datasets llama_index llama-index-llms-cohere llama-index-embeddings-cohere

1
import cohere
2
import datasets
3
from llama_index.core import StorageContext, VectorStoreIndex, load_index_from_storage
4
from llama_index.core.schema import TextNode
5
from llama_index.embeddings.cohere import CohereEmbedding
6
import pandas as pd
7
8
import json
9
from pathlib import Path
10
from tqdm import tqdm
11
from typing import List

1
api_key = "" # <your api="" key="">
2
co = cohere.Client(api_key=api_key)

1
data = datasets.load_dataset("sauravjoshi23/aws-documentation-chunked")
2
print(data)
3
4
map_id2index = {sample["id"]: index for index, sample in enumerate(data["train"])}

/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:88: UserWarning:
The secret `HF_TOKEN` does not exist in your Colab secrets.
To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.
You will be able to reuse this secret in all of your notebooks.
Please note that authentication is recommended but still optional to access public models or datasets.
    warnings.warn(
DatasetDict({
    train: Dataset({
        features: ['id', 'text', 'source'],
        num_rows: 187147
    })
})

1
overwrite = True # only compute index if it doesn't exist
2
path_index = Path(".") / "aws-documentation_index_cohere"
3
4
embed_model = CohereEmbedding(
5
    cohere_api_key=api_key,
6
    model_name="embed-v4.0",
7
)
8
9
if not path_index.exists() or overwrite:
10
    # Documents are prechunked. Keep them as-is for now
11
    stub_len = len("https://github.com/siagholami/aws-documentation/tree/main/documents/")
12
    documents = [
13
        # -- for indexing full dataset --
14
        TextNode(
15
            text=sample["text"],
16
            title=sample["source"][stub_len:], # save source minus stub
17
            id_=sample["id"],
18
        ) for sample in data["train"]
19
        # -- for testing on subset --
20
        # TextNode(
21
        #     text=data["train"][index]["text"],
22
        #     title=data["train"][index]["source"][stub_len:],
23
        #     id_=data["train"][index]["id"],
24
        # ) for index in range(1_000)
25
    ]
26
    index = VectorStoreIndex(documents, embed_model=embed_model)
27
    index.storage_context.persist(path_index)
28
29
else:
30
    storage_context = StorageContext.from_defaults(persist_dir=path_index)
31
    index = load_index_from_storage(storage_context, embed_model=embed_model)

1
retriever = index.as_retriever(similarity_top_k=top_k)

1
class RetrieverWithRerank:
2
    def __init__(self, retriever, api_key):
3
        self.retriever = retriever
4
        self.co = cohere.Client(api_key=api_key)
5
6
    def retrieve(self, query: str, top_n: int):
7
        # First call to the retriever fetches the closest indices
8
        nodes = self.retriever.retrieve(query)
9
        nodes = [
10
            {
11
                "text": node.node.text,
12
                "llamaindex_id": node.node.id_,
13
            }
14
            for node
15
            in nodes
16
        ]
17
        # Call co.rerank to improve the relevance of retrieved documents
18
        reranked = self.co.rerank(query=query, documents=nodes, model="rerank-english-v3.0", top_n=top_n)
19
        nodes = [nodes[node.index] for node in reranked.results]
20
        return nodes
21
22
23
top_k = 60 # how many documents to fetch on first pass
24
top_n = 20 # how many documents to sub-select with rerank
25
26
retriever = RetrieverWithRerank(
27
    index.as_retriever(similarity_top_k=top_k),
28
    api_key=api_key,
29
)

1
query = "What happens to my Amazon EC2 instances if I delete my Auto Scaling group?"
2
3
documents = retriever.retrieve(query, top_n=top_n)
4
5
resp = co.chat(message=query, model="command-r-08-2024", temperature=0., documents=documents)
6
print(resp.text)

1
def build_answer_with_citations(response):
2
    """ """
3
    text = response.text
4
    citations = response.citations
5
6
    # Construct text_with_citations adding citation spans as we iterate through citations
7
    end = 0
8
    text_with_citations = ""
9
10
    for citation in citations:
11
        # Add snippet between last citatiton and current citation
12
        start = citation.start
13
        text_with_citations += text[end : start]
14
        end = citation.end  # overwrite
15
        citation_blocks = " [" + ", ".join([stub[4:] for stub in citation.document_ids]) + "] "
16
        text_with_citations += text[start : end] + citation_blocks
17
    # Add any left-over
18
    text_with_citations += text[end:]
19
20
    return text_with_citations
21
22
grounded_answer = build_answer_with_citations(resp)
23
print(grounded_answer)

1
url = "https://github.com/siagholami/aws-documentation/blob/main/QA_true.csv?raw=true"
2
qa_pairs = pd.read_csv(url)
3
qa_pairs.sample(2)

1
LLM_EVAL_TEMPLATE = """## References
2
{references}
3
4
QUESTION: based on the above reference documents, answer the following question: {question}
5
ANSWER: {answer}
6
STUDENT RESPONSE: {completion}
7
8
Based on the question and answer above, grade the studen't reponse. A correct response will contain exactly \
9
the same information as in the answer, even if it is worded differently. If the student's reponse is correct, \
10
give it a score of 1. Otherwise, give it a score of 0. Let's think step by step. Return your answer as \
11
as a compilable JSON with the following structure:
12
{{
13
    "reasoning": <reasoning>,
14
    "score: <score 0="" 1="" of="" or="">,
15
}}"""
16
17
18
def get_rank_of_golden_within_retrieved(golden: str, retrieved: List[dict]) -> int:
19
    """
20
    Returns the rank that the golden document (single) has within the retrieved documents
21
    * `golden` contains the source of the document, e.g. 'amazon-ec2-user-guide/EBSEncryption.md'
22
    * `retrieved` has a list of responses with key 'llamaindex_id', which links back to document sources
23
    """
24
    # Create {document: rank} map using llamaindex_id (count first occurrence of any document; they can
25
    # appear multiple times because they're chunked)
26
    doc_to_rank = {}
27
    for rank, doc in enumerate(retrieved):
28
        # retrieve source of document
29
        _id = doc["llamaindex_id"]
30
        source = data["train"][map_id2index[_id]]["source"]
31
        # format as in dataset
32
        source = source[stub_len:]  # remove stub
33
        source = source.replace("/doc_source", "")  # remove /doc_source/
34
        if source not in doc_to_rank:
35
            doc_to_rank[source] = rank + 1
36
37
    # Return rank of `golden`, defaulting to len(retrieved) + 1 if it's absent
38
    return doc_to_rank.get(golden, len(retrieved) + 1)

1
from tqdm import tqdm
2
3
answers = []
4
golden_answers = []
5
ranks = []
6
grading_prompts = []  # best computed in batch
7
8
for _, row in tqdm(qa_pairs.iterrows(), total=len(qa_pairs)):
9
    query, golden_answer, golden_doc = row["Question"], row["Answer_True"], row["Document_True"]
10
    golden_answers.append(golden_answer)
11
12
    # --- Produce answer using retriever ---
13
    documents = retriever.retrieve(query, top_n=top_n)
14
    resp = co.chat(message=query, model="command-r-08-2024", temperature=0., documents=documents)
15
    answer = resp.text
16
    answers.append(answer)
17
18
    # --- Do some prework for evaluation later ---
19
    # Rank
20
    rank = get_rank_of_golden_within_retrieved(golden_doc, documents)
21
    ranks.append(rank)
22
    # Score: construct the grading prompts for LLM evals, then evaluate in batch
23
    # Need to reformat documents slightly
24
    documents = [{"index": str(i), "text": doc["text"]} for i, doc in enumerate(documents)]
25
    references_text = "\n\n".join("\n".join([f"{k}: {v}" for k, v in doc.items()]) for doc in documents)
26
    # ^ snippet looks complicated, but all it does it unpack all kwargs from `documents`
27
    # into text separated by \n\n
28
    grading_prompt = LLM_EVAL_TEMPLATE.format(
29
        references=references_text, question=query, answer=golden_answer, completion=answer,
30
    )
31
    grading_prompts.append(grading_prompt)

1
results = pd.DataFrame()
2
results["answer"] = answers
3
results["golden_answer"] = qa_pairs["Answer_True"]
4
results["rank"] = ranks

1
scores = []
2
reasonings = []
3
4
def remove_backticks(text: str) -> str:
5
  """
6
  Some models are trained to output JSON in Markdown formatting:
7
  ```json {json object}```
8
  Remove the backticks from those model responses so that they become
9
  parasable by json.loads.
10
  """
11
  if text.startswith("```json"):
12
      text = text[7:]
13
  if text.endswith("```"):
14
      text = text[:-3]
15
  return text
16
17
18
for prompt in tqdm(grading_prompts, total=len(grading_prompts)):
19
  resp = co.chat(message=prompt, model="command-a-03-2025", temperature=0.)
20
  # Convert response to JSON to extract the `score` and `reasoning` fields
21
  # We remove backticks for compatibility with different LLMs
22
  parsed = json.loads(remove_backticks(resp.text))
23
  scores.append(parsed["score"])
24
  reasonings.append(parsed["reasoning"])

1
results["score"] = scores
2
results["reasoning"] = reasonings

1
print(f"Average score: {results['score'].mean():.3f}")

1
import matplotlib.pyplot as plt
2
import seaborn as sns
3
4
sns.set_theme(style="darkgrid", rc={"grid.color": ".8"})
5
6
results["rank_shifted_left"] = results["rank"] - 0.1
7
results["rank_shifted_right"] = results["rank"] + 0.1
8
9
f, ax = plt.subplots(figsize=(5, 3))
10
sns.histplot(data=results.loc[results["score"] == 1], x="rank_shifted_left", color="skyblue", label="Correct answer", binwidth=1)
11
sns.histplot(data=results.loc[results["score"] == 0], x="rank_shifted_right", color="red", label="False answer", binwidth=1)
12
13
ax.set_xticks([1, 5, 0, 10, 15, 20])
14
ax.set_title("Rank of golden document (max means golden doc. wasn't retrieved)")
15
ax.set_xlabel("Rank")
16
ax.legend();