RAG with Bedrock

Last updated on 2026-11-02 | Edit this page

Estimated time: 50 minutes

Overview

Questions

  • How do we build a complete RAG pipeline on AWS without provisioning a GPU?
  • How do we turn a collection of PDFs into retrievable, citable chunks stored in S3?
  • How do we call Bedrock-hosted embedding and generation models so that the spend is tagged to our project?

Objectives

  • Read the WattBot metadata, questions and PDF corpus from the shared bucket, and cache what you derive from them (chunks.jsonl) in your own bucket for the later episodes.
  • Convert PDFs into page-level records and overlapping text chunks.
  • Create tagged application inference profiles so that every Bedrock call is attributed to your project.
  • Embed chunks with Amazon Titan Text Embeddings and retrieve relevant context with cosine similarity.
  • Generate WattBot-format answers and explanations with a Bedrock-hosted model (DeepSeek V3) and score them against train_QA.csv.
  • Build a second index from page images with a multimodal embedding model and measure what it changes for the table and figure questions.

Open the notebook from: /ML_with_AWS_SageMaker/notebooks/11_RAG-Bedrock.ipynb. Select the conda_pytorch kernel. This is the first hands-on RAG episode, and it uses Amazon Bedrock, the default the overview recommended. We build the whole RAG pipeline for the WattBot 2026 challenge from a small CPU-only SageMaker notebook. Nothing here needs a GPU, because both the text embeddings and the answer generation run on Bedrock-hosted models, which bill per token and leave nothing running between calls. The two models are Amazon Titan Text Embeddings V2, which turns WattBot chunks into vectors, and DeepSeek V3, which writes the answers and explanations.

This episode also prepares the corpus. It unpacks the WattBot PDFs, splits them into pages and overlapping chunks, and caches the chunks in S3. The Processing-jobs episode and the optional GPU-notebook episode reuse those chunks and run the same retrieval, evaluation and scoring on self-hosted models, so cost, latency and effort can be compared directly.

Why Bedrock for WattBot?

With a self-hosted model, as in the Processing-jobs episode, you choose the model, keep its version, and make sure the instance has the GPU memory for it. That is fine for an experiment and a burden once several teams want the same pipeline. Running the embedding and generation steps on Bedrock means:

  • Managed models. Models from Anthropic, Amazon, DeepSeek and others, with no container images or CUDA versions of your own.
  • Paying per token. Instead of a GPU instance that may sit idle, you pay for the input and output tokens of each call. For most workloads that is cheaper; for a large offline batch with a small model, a GPU job can still win.
  • One set of models for a group. A lab or a course can agree on a few Bedrock models more easily than on a fleet of GPU instances.

In this notebook both the embedding step and the answer and explanation steps run on Bedrock. The embeddings are saved as embeddings_bedrock.npy, so they can be compared with the self-hosted run in the Processing-jobs episode.

Models used in this episode

Both models are Bedrock-hosted foundation models:

  • Amazon Titan Text Embeddings V2 (amazon.titan-embed-text-v2:0): a general text-embedding model for retrieval, clustering and classification. It returns vectors of 256, 512 or 1,024 dimensions, with presets for retrieval. AWS does not publish its parameter count.
  • DeepSeek V3 (deepseek.v3-v1:0): an open-weight mixture-of-experts model, 671B parameters in total and about 37B active per token, that Bedrock serves on demand, so there is no GPU to manage. Its per-token price is well below the flagship models’ (the Bedrock pricing page has current rates), which matters when there is one call per question. The same code works for any chat model in the catalog, because Bedrock’s one chat call, converse, takes the same input for every provider, a system prompt and a list of messages, and returns the reply in the same shape, so changing base_generation_model_id to a Claude, Llama or Mistral ID changes nothing else.
  • Multimodal models for tables and figures: Bedrock also serves models that embed or read images. Step 13 uses one, Amazon Nova Multimodal Embeddings, on the table and figure questions.

The model catalog in the Bedrock console lists every model available in the region (Claude variants, Amazon models and partner models), each with a model card and a link to its pricing.

Callout

Choosing the two models

The catalog tells you what exists and what it costs; it does not tell you what is good. Three places do, and all three are free to read:

  • Embedding models: the MTEB leaderboard scores embedding models on retrieval benchmarks. Filter to the retrieval task and your language; the top of the table changes every few months, and Titan V2 is a sound middle-of-the-pack choice that Bedrock serves without setup.
  • Generation models: LMArena ranks chat models by blind human preference, and Artificial Analysis plots quality against price per million tokens, the chart to read for RAG, since the answer is written from retrieved text, so a mid-priced model often reads as well as a flagship at a tenth of the cost.
  • Your own questions: the benchmarks rank models on other people’s data. train_QA.csv has labeled questions for this corpus, and the evaluation at the end of this episode scores a model on them and prices the run from its token counts. Changing base_generation_model_id and rerunning is the whole experiment. A model that wins on a leaderboard and loses on your questions has told you something the leaderboard cannot.

The habit is the same as for instances. Pick the cheapest candidate that scores well enough, measure it on your data, and move up only when the measurement says to.

Notebook + dataset setup

For this episode you need a SageMaker notebook instance, but not a GPU one. A small CPU instance such as ml.t3.medium is enough, because the models run on Bedrock, not in the notebook. The notebook’s execution role must allow Bedrock access, which a default SageMaker role does not. On the class account both roles the console offers carry it. On your own account, attach the Bedrock policy from the Account preparation checklist instructor page to your role before starting; the first Bedrock call fails with an access-denied error naming the role until it is there.

See Instances for ML for further guidance on instance types.

Two buckets, no uploads

This episode reads from two buckets and you create neither of them here:

  • The shared bucket, ml-marathon-shared, already set in the setup cells below, holds the WattBot inputs under wattbot/: metadata.csv (the index of papers), train_QA.csv (labeled questions and answers), corpus.zip (the PDFs) and pages.jsonl (the text of every PDF page, extracted once). It is read-only and there is one copy for the whole class.
  • Your bucket from the S3 episode (challenge-lastnamefirstname) is where this episode writes what it derives, under a wattbot/ prefix: the chunked text (wattbot/chunks.jsonl), the embeddings (wattbot/embeddings/) and the predictions (wattbot/solutions/). Its tags cover its storage cost; your notebook instance carries the tags you gave it when you created it, and in this episode those two resources plus the per-token Bedrock calls (tagged through an inference profile, see “Cost tracking” below) are the only things being billed.
Callout

Working through the lesson on your own?

If there is no instructor-staged bucket, stage the four input files in your own bucket under the same keys and point both variables at it: metadata.csv and train_QA.csv are in data/data.zip in the lesson repository, corpus.zip is built once by scripts/build_wattbot_corpus.py, and instructors/shared-data-bucket.md has the upload commands. Then set shared_bucket = bucket_name in the setup cell below and everything else works unchanged.

Step 1 – Configure the notebook

Three things are new here compared with the training episodes, so the setup is split into four cells. The session, region and S3 client are as before, plus the execution role, printed because the Bedrock calls below run as it.

PYTHON 

import os
import json
import time
from typing import Dict, Any, List
from pathlib import Path

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError
import pandas as pd
import numpy as np
import sys
sys.path.insert(0, "/home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts")   # the lesson's helper module, cloned in the controllers episode
import helpers
import sagemaker

session = sagemaker.Session()
region = session.boto_region_name
s3 = boto3.client("s3", region_name=region)
role = sagemaker.get_execution_role()
print("Role:", role)   # the execution role: on the class account it carries the Bedrock permissions; on your own, it must

The models. RAG needs two, an embedding model, which turns a piece of text into a vector so that similar passages land near each other, and a generation model, the chat model that writes the answer from the passages it is given. Both are catalog models on Bedrock, named by a model ID; the model catalog in the console lists everything available in the region with its ID; prices are on the pricing page, and the supported models page in the documentation is the same list as a table. Prices are per million tokens on the Bedrock pricing page. In October 2026, DeepSeek V3 on demand is $0.62 in and $1.85 out, and Titan Text Embeddings V2 is $0.02 in, which is why the whole WattBot run below costs cents. We write the IDs down now but will not call them by these IDs directly, because the cost-tracking step below wraps each one in a tagged inference profile first, because that is the only way Bedrock spend gets attributed to a project.

PYTHON 

# Base model IDs. Any chat model in the Bedrock catalog works as the generator (a Claude or Llama ID, say),
# because converse, Bedrock's chat call, takes the same messages and returns the same shape for every provider.
base_generation_model_id = "deepseek.v3-v1:0"             # writes the answers; $0.62 in / $1.85 out per million tokens (Oct 2026)
base_embedding_model_id = "amazon.titan-embed-text-v2:0"  # turns text into vectors for retrieval; $0.02 per million tokens

The files. The inputs in the shared bucket are the paper list, the questions with their answers, and the zipped corpus of PDFs. Everything this episode derives from them goes to your own bucket, under the same wattbot/ prefix, which will hold the chunked text, the chunk embeddings, and the answers the pipeline produces.

PYTHON 

!mkdir -p /home/ec2-user/SageMaker/wattbot   # a project folder beside the lesson clone, so data and outputs stay out of the repository
%cd /home/ec2-user/SageMaker/wattbot

shared_bucket = "ml-marathon-shared"   # the instructor-staged, read-only bucket with the data and model weights; change only on your own account
bucket_name = "sinkorswim-doejohn"     # <-- EDIT ME: your own bucket from the S3 episode

# inputs, in the shared bucket
metadata_key = "wattbot/metadata.csv"      # one row per paper: id, title, url
train_key = "wattbot/train_QA.csv"         # the questions, with reference answers for scoring
corpus_key = "wattbot/corpus.zip"          # the PDFs, one per paper

# derived, written to YOUR bucket
chunks_key = "wattbot/chunks.jsonl"                                # the corpus cut into passages (end of the chunking step)
embeddings_key = "wattbot/embeddings/embeddings_bedrock.npy"       # one vector per chunk
solutions_key = "wattbot/solutions/wattbot_solutions_bedrock.csv"  # the pipeline's answers, in the competition's format

local_data_dir = "./data"              # working copies on this notebook instance
os.makedirs(local_data_dir, exist_ok=True)

The inputs, as they sit in the shared bucket. The three files named above are at the top of wattbot/; the rest of the folder is what the instructor staged for the caching steps later in the episode.

PYTHON 

helpers.show_bucket(shared_bucket, "wattbot/")   # every file under the prefix, with its size

Two Bedrock clients. Bedrock has two APIs. bedrock-runtime is the one that calls models (converse, the chat call, for answers; invoke_model, the raw call with a provider-specific body, for embeddings); it is what every answer and every embedding in this episode goes through. bedrock is the control plane, for listing models and creating the inference profiles that carry our tags. The runtime client gets a retry setting because on-demand Bedrock throughput is rate-limited per model, so a loop of a few thousand embedding calls will occasionally be throttled, and the adaptive mode backs off and retries instead of raising.

PYTHON 

bedrock_runtime = boto3.client(
    "bedrock-runtime",
    region_name=region,
    config=Config(retries={"max_attempts": 10, "mode": "adaptive"}),   # back off and retry when throttled
)
bedrock = boto3.client("bedrock", region_name=region)                  # control plane: list models, create inference profiles

Step 2 – Get the shared WattBot helpers

The retrieval, prompting, answer parsing and scoring logic of the WattBot pipeline is the same whether the models run on Bedrock, in a Processing job or on a notebook GPU. It lives once in the lesson repository, in scripts/wattbot_rag.py, and every RAG episode imports it; the only code that differs per episode is the two model calls. Copy the module from the lesson clone to the folder next to this notebook (the Processing Jobs episode ships the same file inside its jobs).

PYTHON 

!cp /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/wattbot_rag.py /home/ec2-user/SageMaker/ML_with_AWS_SageMaker/scripts/Score.py .

PYTHON 

from wattbot_rag import (
    ANSWER_SYSTEM_PROMPT,
    build_docid_to_url,
    compute_wattbot_score,
    retrieve_top_k,
    split_ref_ids,
    run_single_qa,
)

Open wattbot_rag.py in the Jupyter file browser as you go. Each helper is short.

Step 3 – Load the metadata and training questions

Pull the two CSV files from the shared bucket into the notebook and inspect them. Two helpers come first. download_from_s3 copies one object from the shared bucket (or any bucket you name) into the local data folder, and smart_read_csv reads a CSV by trying a few encodings in turn, because author names and titles in metadata.csv are not always valid UTF-8.

PYTHON 

import zipfile


def download_from_s3(key: str, local_name: str, bucket: str = shared_bucket) -> str:
    """Download a file from S3 (the shared bucket by default) to local_data_dir."""
    local_path = os.path.join(local_data_dir, local_name)
    print(f"Downloading s3://{bucket}/{key} -> {local_path}")
    s3.download_file(bucket, key, local_path)
    return local_path


def smart_read_csv(path: str) -> pd.DataFrame:
    """Read a CSV, trying utf-8 and then the common Windows and Latin encodings."""
    encodings = ["utf-8", "latin1", "ISO-8859-1", "cp1252"]
    last_error = None
    for enc in encodings:
        try:
            return pd.read_csv(path, encoding=enc)
        except Exception as e:
            last_error = e
    if last_error is not None:
        raise last_error
    raise RuntimeError(f"Unable to read CSV at {path}")

Now the two files. train_QA.csv is the labeled questions and answers we score against, metadata.csv is one row per paper with its id, title and URL.

PYTHON 

train_qa_path = download_from_s3(train_key, "train_QA.csv")      # labeled questions + answers: what we score against
metadata_path = download_from_s3(metadata_key, "metadata.csv")   # one row per paper (id, title, url): the corpus index

train_df = smart_read_csv(train_qa_path)
metadata_df = smart_read_csv(metadata_path)

print("train_QA.csv columns:", train_df.columns.tolist())
print("metadata.csv columns:", metadata_df.columns.tolist())
print("\nNumber of training QAs:", len(train_df))
print("Number of metadata rows:", len(metadata_df))

train_df.head(15)

Step 4 – Get the PDF corpus

The corpus is the set of PDFs listed in metadata.csv, one per paper, named <id>.pdf. Downloading them from the publishers’ sites is slow and unreliable when fifty notebooks do it at once, so the instructors did it once with scripts/build_wattbot_corpus.py and staged the result as corpus.zip in the shared bucket. We download that one file, unpack it, and delete the zip, because the instance’s volume is 20 GB, and the zip plus the unpacked PDFs are close to 700 MB.

PYTHON 

PDF_DIR = os.path.join(local_data_dir, "pdfs")
os.makedirs(PDF_DIR, exist_ok=True)

corpus_path = download_from_s3(corpus_key, "corpus.zip")
with zipfile.ZipFile(corpus_path, "r") as zf:
    zf.extractall(PDF_DIR)
os.remove(corpus_path)   # the zip is 270 MB and the unpacked PDFs are 400 MB; the notebook disk is small

pdf_files = sorted(f for f in os.listdir(PDF_DIR) if f.lower().endswith(".pdf"))
print(f"{len(pdf_files)} PDFs in {PDF_DIR}; metadata lists {len(metadata_df)} papers")
Callout

Where the corpus came from

metadata.csv has a url column pointing at each paper. scripts/build_wattbot_corpus.py in the lesson repository walks that column, downloads every PDF, zips them, and then extracts the text of every page into pages.jsonl; it reports any URL that failed so the instructor can decide what to do about it. If you are working on your own, run it once and upload both files to your bucket under wattbot/ (with shared_bucket = bucket_name, as in the callout above). That is the whole difference between “the data” and “a copy of the data”. The corpus is built once and read many times.

Step 5 – Turn PDFs into page-level “documents”

The next unit of work is a list of page-level records, one per PDF page. The instructor’s staging script already built it for this corpus (pages.jsonl in the shared bucket), so the cell below loads the staged result and falls back to extracting only if it is missing. Each record stores:

  • text: page text (as extracted by pypdf).
  • doc_id: short ID from metadata.csv (e.g., strubell2019).
  • title: title of the document.
  • url: original PDF URL.
  • page_num: zero-based page index.
  • page_label: label used inside the PDF (often 1-based).

Later, we will chunk these pages into smaller overlapping segments for embedding.

Why chunks keep their page

Every chunk carries the page it came from, so the document id, URL and page label stay attached to the text, a retrieved chunk can be traced to the page the model saw, the text can be cleaned at the page level before chunking, and chunk sizes can be changed later without extracting the PDFs again.

PYTHON 

%pip install -q pypdf

pdfs_to_page_docs in wattbot_rag.py opens each PDF listed in the metadata with pypdf and turns every page into one record holding its text, the document id, title and URL, and the page number. It is forty lines of plain Python with nothing AWS in it; open the file to read it. On a CPU notebook it takes several minutes over this corpus, so the cell first looks for pages.jsonl, the same records written once when the corpus was staged, and runs the extraction only if that file is not there.

PYTHON 

from wattbot_rag import pdfs_to_page_docs   # the extraction: open each PDF with pypdf, one record per page; in scripts/wattbot_rag.py

start = time.time()
try:
    pages_path = download_from_s3("wattbot/pages.jsonl", "pages.jsonl")   # staged next to corpus.zip by the corpus build script
    page_docs = [json.loads(line) for line in open(pages_path)]
    print(f"Loaded {len(page_docs)} pre-extracted page records in {time.time() - start:.1f} s")
except ClientError:
    page_docs = pdfs_to_page_docs(metadata_df, PDF_DIR)
    print(f"Extracted {len(page_docs)} page-level records from {len(metadata_df)} PDFs in {(time.time() - start) / 60:.1f} min")
page_docs[0] if page_docs else None

A paper whose PDF the corpus build could not download has no pages here, and a question that cites it can only score zero, which would measure the download, not the pipeline. drop_questions_without_papers removes those questions from the evaluation set and says which papers were missing; the unanswerable questions, which cite nothing, stay.

PYTHON 

from wattbot_rag import drop_questions_without_papers

train_df = drop_questions_without_papers(train_df, {p["doc_id"] for p in page_docs})
print(f"{len(train_df)} questions can be answered from this corpus")

Step 6 – Text chunking

Pages are split into chunks because a page is the wrong size for retrieval: too long, and one vector has to stand for several topics at once; too short, and a chunk carries no context. The chunker here is a sliding window over characters:

  • chunk_size_chars: the most characters in a chunk (1,000 to 1,500 is usual).
  • chunk_overlap_chars: how many characters consecutive chunks share (200 here), so a sentence cut by a boundary is whole in one of them.

In your own work, splitting on headings, sections or sentence boundaries, called semantic chunking, is the usual next step.

PYTHON 

def split_text_into_chunks(
    text: str,
    chunk_size_chars: int = 1200,     # max characters per chunk (~300 tokens)
    chunk_overlap_chars: int = 200,   # consecutive chunks share this much, so a sentence cut at the edge survives
) -> List[str]:
    """Split `text` into overlapping character-based chunks.

    This is a simple baseline; more advanced versions might:
    - split on sentence boundaries, or
    - merge short paragraphs and respect section headings.
    """
    text = text.strip()
    if not text:
        return []

    chunks: List[str] = []
    start = 0
    text_len = len(text)

    while start < text_len:
        end = min(start + chunk_size_chars, text_len)
        chunk = text[start:end]
        chunks.append(chunk)
        if end == text_len:
            break
        # Move the window forward, keeping some overlap
        start = end - chunk_overlap_chars

    return chunks


def make_chunked_docs(
    page_docs: List[Dict[str, Any]],
    chunk_size_chars: int = 1200,
    chunk_overlap_chars: int = 200,
) -> List[Dict[str, Any]]:
    """Turn page-level records into smaller overlapping text chunks.

    Each chunk keeps a pointer back to its document and page metadata.
    """
    chunked: List[Dict[str, Any]] = []
    for page in page_docs:
        page_text = page["text"]
        chunks = split_text_into_chunks(
            page_text,
            chunk_size_chars=chunk_size_chars,
            chunk_overlap_chars=chunk_overlap_chars,
        )
        for idx, chunk_text in enumerate(chunks):
            chunked.append(
                {
                    "text": chunk_text,
                    "doc_id": page["doc_id"],
                    "title": page["title"],
                    "url": page["url"],
                    "page_num": page["page_num"],
                    "page_label": page["page_label"],
                    "total_pages": page["total_pages"],
                    "chunk_idx_in_page": idx,   # position within the page; the rest is copied from the page record
                }
            )
    return chunked

Chunking is cheap CPU work, but there is no reason to repeat it on every run. We save the chunks locally as data/chunks.jsonl (one JSON record per line) and reload the file if it already exists.

PYTHON 

chunks_jsonl_path = os.path.join(local_data_dir, "chunks.jsonl")

def save_chunked_docs_jsonl(path, chunks):
    with open(path, "w", encoding="utf-8") as f:
        for rec in chunks:
            json.dump(rec, f, ensure_ascii=False)
            f.write("\n")


def load_chunked_docs_jsonl(path):
    with open(path, "r", encoding="utf-8") as f:
        return [json.loads(line) for line in f]

# -------------------------------------------------------------------
# Cached chunking logic
# -------------------------------------------------------------------
if os.path.exists(chunks_jsonl_path):
    print(f"Found existing chunk file: {chunks_jsonl_path}")
    chunked_docs = load_chunked_docs_jsonl(chunks_jsonl_path)
    print("Loaded chunked docs:", len(chunked_docs))
else:
    print("No chunk file found. Running chunking step...")
    chunked_docs = make_chunked_docs(page_docs)
    save_chunked_docs_jsonl(chunks_jsonl_path, chunked_docs)
    print(f"Saved chunked docs to {chunks_jsonl_path}")

# Show first chunk
print("Raw pages:", len(page_docs))
print("Chunked docs:", len(chunked_docs))
chunked_docs[0] if chunked_docs else None

PYTHON 

# Upload to YOUR bucket so the Processing Jobs episode (and future runs) can reuse the
# chunks without repeating the unpacking and chunking steps.
print(f"Uploading chunked docs to s3://{bucket_name}/{chunks_key} ...")
s3.upload_file(chunks_jsonl_path, bucket_name, chunks_key)
print("Upload complete.")

Step 7 – Cost tracking through a tagged profile

Every SageMaker job in this lesson is tagged at launch, because tags never propagate from the notebook. Bedrock is stricter still. An on-demand invoke_model or converse call has no tags parameter at all. If you call deepseek.v3-v1:0 or amazon.titan-embed-text-v2:0 by its model ID, the usage lands in the account’s Bedrock bill as anonymous per-token charges. The notebook’s tags, the bucket’s tags, and the tags you set on anything else are irrelevant to it. On a shared account that means nobody can tell which team spent what.

The fix is not a tag on the call; it is a different thing to call. An application inference profile is a small resource you create in your own account, an alias that points at one base model and carries your three tags. It has its own ARN, and that ARN goes wherever a model ID goes. Nothing else changes. Same model, same price per token, same answers. The only difference is in the bill, where usage made through the alias shows up with the alias’s tags.

Before and after, in the generation call:

modelId you pass Who it is billed to
Raw model ID deepseek.v3-v1:0 nobody: anonymous DeepSeek spend on the account
Your profile arn:aws:bedrock:us-east-2:<account-id>:application-inference-profile/adewnu8wqxpy Name=DoeJohn, Project=SinkOrSwim, Purpose=Inference

So cost tracking on Bedrock comes to three things. Create the alias once per model (one for DeepSeek, one for Titan, because an alias points at a single model), keep its ARN in a variable, and use that variable as modelId from then on. The profile is free; you pay only for tokens, as before.

One call per model does all of it. tagged_profile_arn in wattbot_rag.py looks for a profile named after your project and user name, creates it with the three tags if it is not there, and returns its ARN. One tagged profile per model, its ARN in place of the model ID.

PYTHON 

from wattbot_rag import tagged_profile_arn

name = "DoeJohn"         # replace with your name
project = "SinkOrSwim"   # your challenge: WattBot2026, BadgerScribe, BadgerCode, BrainDecoding, PoisonedWell or SnapshotWI; SinkOrSwim is the workshop itself
print(f"Tags on the profiles, and so on every Bedrock call: Name={name}, Project={project}, Purpose=Inference")   # check them first: a profile tagged DoeJohn bills as DoeJohn

# From here on, these ARNs are what we pass as modelId. The base IDs are never sent to the runtime again.
bedrock_model_id = tagged_profile_arn(bedrock, base_generation_model_id, "wattbot-gen", name, project)
embedding_model_id_bedrock = tagged_profile_arn(bedrock, base_embedding_model_id, "wattbot-embed", name, project)
Callout

The same thing from a terminal

Outside a notebook it is one CLI command per model, which is all the helper wraps:

It prints the profile ARN; pass that as modelId. The helper adds only the reuse check, so rerunning does not create duplicates.

Callout

Checking the profiles

  • You can confirm the routing with bedrock.list_inference_profiles(typeEquals="APPLICATION"); the console’s inference-profile pages list only the system-defined profiles.
  • The keys must be the three approved ones, Name, Project and Purpose, with that spelling and capitalization; any other key is ignored by the bill (the S3 episode says why).
  • The profile itself has no hourly or storage charge. You only pay per token, as before.
Challenge

Who is paying for these tokens?

List the application inference profiles in the account and find your two. For each, print its name, the model it points at, and its tags (bedrock.list_tags_for_resource takes the profile ARN). Then answer this. If a teammate skipped this step and called deepseek.v3-v1:0 by its raw model ID for the rest of the episode, what would Cost Explorer show when the organizers filter the month’s Bedrock spend by Project?

PYTHON 

profiles = bedrock.list_inference_profiles(typeEquals="APPLICATION")["inferenceProfileSummaries"]
print(f"{len(profiles)} application inference profiles in this account/region")
for p in profiles:
    if p["inferenceProfileArn"] in {bedrock_model_id, embedding_model_id_bedrock}:
        tags = bedrock.list_tags_for_resource(resourceARN=p["inferenceProfileArn"])["tags"]
        print(p["inferenceProfileName"], "->", [m["modelArn"].split("/")[-1] for m in p["models"]])
        print("   tags:", {t["key"]: t["value"] for t in tags})

Filtering by Project shows only the usage that went through a tagged profile. The teammate’s calls are still on the bill, under the Bedrock service line for deepseek.v3-v1:0, but with no Project, Name or Purpose value, so in the grouped view they fall into the “No tag key: Project” bucket, together with every other untagged call in the account. There is no way to attribute them after the fact; tags are applied when the usage is metered, not later. That is why the base model ID is never passed to the runtime in this episode.

Step 8 – Embed the chunks with Titan

Everything the RAG loop needs is now in memory, the chunks, the questions, and a way to resolve a document id to its URL for the ref_url column.

PYTHON 

print(f"Chunks: {len(chunked_docs)}")
print(f"Train QAs: {len(train_df)}")
print(f"Metadata rows: {len(metadata_df)}")

docid_to_url = build_docid_to_url(metadata_df)   # doc_id -> url, for the ref_url column of the submission
print(f"docid_to_url has {len(docid_to_url)} entries.")

Titan Text Embeddings V2 takes one string per invoke_model call and returns a 1,024-dimensional normalized vector. embedding_model_id_bedrock is the ARN of our tagged profile for it, so every call is attributed to our Project/Name/Purpose tags. The next three cells set where the result is saved, define the call that embeds one string, and run the loop that embeds every chunk once and caches the matrix. First the paths.

PYTHON 

data_dir = Path(local_data_dir)
emb_save_path = data_dir / "embeddings_bedrock.npy"             # the finished matrix, one row per chunk
emb_partial_path = data_dir / "embeddings_bedrock.partial.npy"  # progress file, so an interrupted run resumes

The call itself. invoke_model takes a provider-specific JSON body; Titan’s is {"inputText": ...}, and its response carries the vector under embedding. embed_query is the same call wrapped for retrieval time, when we embed one question at a time.

PYTHON 

def bedrock_embed_text(text: str, model_id: str = embedding_model_id_bedrock) -> List[float]:
    """Call the Titan embedding model (through our tagged profile) for one string."""
    response = bedrock_runtime.invoke_model(
        modelId=model_id,
        contentType="application/json",
        accept="application/json",
        body=json.dumps({"inputText": text}),   # Titan's own request shape; add "dimensions": 256 or 512 to shrink the 1,024 default
    )
    response_body = json.loads(response["body"].read())
    embedding = response_body.get("embedding")   # a list of floats, one per dimension
    if embedding is None:
        raise ValueError(f"No 'embedding' found in response: {response_body}")
    return embedding


def embed_query(question: str) -> np.ndarray:
    """The retrieval-time half of the pipeline: the same model, one question at a time."""
    return np.asarray(bedrock_embed_text(question), dtype="float32")

Embedding the corpus, about 9,000 chunks, is ten minutes and a few cents, so do not do it twice. The cell after next looks for the finished matrix locally, then in the shared bucket if use_instructor_embeddings is True, then in your bucket, and only then computes it, saving progress every 500 chunks so an interrupted run picks up where it stopped.

Callout

In a class, use the instructor’s copy

Bedrock allows this account 300,000 Titan embedding tokens per minute. This corpus is about 2.4 million tokens, so one notebook alone needs ten minutes at the full quota; a room of them shares it, so everyone’s loop slows to a crawl of retries. In class, leave use_instructor_embeddings in the next cell at True. The cache cell then downloads the matrix your instructor computed with the same chunks and the same model from the shared bucket and skips the calls. Working on your own, set it to False and let the loop run; that is the step to practice.

PYTHON 

use_instructor_embeddings = True                            # True in class: download the instructor's matrix. False: embed the 9,000 chunks yourself, ten minutes at the quota
shared_embeddings_key = "wattbot/embeddings_bedrock.npy"   # where the instructor staged it in the shared bucket

One small helper asks S3 whether an object exists, and one checks that a cached matrix belongs to these chunks, since a matrix from an earlier corpus has a different number of rows, and using it would pair every question with the wrong passages.

PYTHON 

def _s3_object_exists(bucket: str, key: str) -> bool:
    try:
        s3.head_object(Bucket=bucket, Key=key)
        return True
    except ClientError:
        return False


def load_embeddings_if_matching(path) -> Any:
    """The matrix at `path` if it has one row per chunk, else None (a matrix from a different corpus)."""
    matrix = np.load(path)
    if matrix.shape[0] == len(chunked_docs):
        return matrix
    print(f"{path} has {matrix.shape[0]} rows but there are {len(chunked_docs)} chunks: a matrix from a different corpus, ignoring it")
    return None


# Cache order: local file, the instructor's copy (in class), your bucket, then recompute (and upload for next time)
chunk_embeddings = None
if emb_save_path.exists():
    chunk_embeddings = load_embeddings_if_matching(emb_save_path)
if chunk_embeddings is None and use_instructor_embeddings and _s3_object_exists(shared_bucket, shared_embeddings_key):
    s3.download_file(shared_bucket, shared_embeddings_key, str(emb_save_path))
    print(f"Downloaded the instructor's embeddings from s3://{shared_bucket}/{shared_embeddings_key}")
    chunk_embeddings = load_embeddings_if_matching(emb_save_path)
if chunk_embeddings is None and _s3_object_exists(bucket_name, embeddings_key):
    s3.download_file(bucket_name, embeddings_key, str(emb_save_path))
    print(f"Downloaded your cached embeddings from s3://{bucket_name}/{embeddings_key}")
    chunk_embeddings = load_embeddings_if_matching(emb_save_path)
if chunk_embeddings is None:
    # Resume from a partial save if a previous run was interrupted
    all_embeddings = list(np.load(emb_partial_path)) if emb_partial_path.exists() else []
    start = len(all_embeddings)
    print(f"Computing embeddings via Bedrock from chunk {start} of {len(chunked_docs)}...")
    t0 = time.time()

    for idx in range(start, len(chunked_docs)):
        all_embeddings.append(bedrock_embed_text(chunked_docs[idx]["text"]))
        if (idx + 1) % 500 == 0 or idx + 1 == len(chunked_docs):
            np.save(emb_partial_path, np.array(all_embeddings, dtype="float32"))
            print(f"Embedded {idx + 1} / {len(chunked_docs)} chunks (progress saved)")

    print(f"Embedded {len(chunked_docs) - start} chunks in {time.time() - t0:.1f} seconds")
    chunk_embeddings = np.array(all_embeddings, dtype="float32")   # the embedding matrix: one row per chunk, one column per dimension
    np.save(emb_save_path, chunk_embeddings)
    emb_partial_path.unlink(missing_ok=True)
    s3.upload_file(str(emb_save_path), bucket_name, embeddings_key)
    print(f"Saved embeddings to {emb_save_path} and s3://{bucket_name}/{embeddings_key}")

assert chunk_embeddings.shape[0] == len(chunked_docs), "embeddings and chunks are out of step; delete the cached files and rerun"
print("Embeddings shape:", chunk_embeddings.shape)

Step 9 – Retrieve with cosine similarity over the embedding matrix

Retrieval is pure NumPy. Embed the question with the same model, take the cosine similarity against every chunk, and keep the top k. retrieve_top_k from wattbot_rag.py does that (open the file to see it, a normalize and a matrix multiply). No vector database is needed for a few thousand chunks; the optional GPU-notebook episode discusses when one becomes worth it.

PYTHON 

first_row = train_df.iloc[0]
print("Sample question:", first_row["question"])

# Cosine similarity of the question vector against every row of the matrix; keep the k best
test_retrieved = retrieve_top_k(embed_query(first_row["question"]), chunk_embeddings, chunked_docs, k=3)
for r in test_retrieved:   # score is the cosine similarity; 1.0 would be an identical vector
    snippet = r["text"].replace("\n", " ")[:160]
    print(f"- score={r['score']:.3f} | doc_id={r['doc_id']} | page={r['page_label']} | {snippet}...")

Step 10 – Generate with DeepSeek V3 through the converse API

Bedrock has two runtime APIs. invoke_model takes a provider-specific request body (which is why the Titan call above builds a {"inputText": ...} body by hand). converse is provider-neutral. It takes a system prompt, a list of messages and an inferenceConfig, in the same shape for DeepSeek, Claude, Llama and the rest. We use converse for generation so the model is a one-line swap. The response carries the generated text and a usage block with the input and output token counts, which is what Bedrock bills on; we keep every usage block so we can price the run at the end.

PYTHON 

bedrock_usage = []   # one {"inputTokens": ..., "outputTokens": ...} per call

def call_bedrock_chat(
    system_prompt: str,
    user_prompt: str,
    max_new_tokens: int = 512,
    model_id: str = bedrock_model_id,   # the tagged profile ARN, never the raw model ID
) -> str:
    """Send one system + user turn to a Bedrock chat model and return the generated text."""
    try:
        response = bedrock_runtime.converse(
            modelId=model_id,
            # system: instructions that apply to the whole exchange
            system=[{"text": system_prompt}],
            # messages: the conversation so far, here a single user turn
            messages=[{"role": "user", "content": [{"text": user_prompt}]}],
            # inferenceConfig: generation settings; temperature 0 makes the output repeatable
            inferenceConfig={"maxTokens": max_new_tokens, "temperature": 0},
        )
    except ClientError as e:
        print(f"ERROR calling Bedrock model {model_id}: {e}")
        raise

    bedrock_usage.append(response["usage"])   # token counts for this call, kept so we can price the run later
    # Some models return extra content blocks (e.g. reasoning); keep only the text ones.
    text_blocks = [block["text"] for block in response["output"]["message"]["content"] if "text" in block]
    return "\n".join(text_blocks).strip()


print(call_bedrock_chat("You are a concise assistant.", "What is 2 + 2? Answer in one short sentence.", max_new_tokens=32))
print("usage of that call:", bedrock_usage[-1])

That function has the signature the shared pipeline expects, generate(system_prompt, user_prompt, max_new_tokens) -> str, and embed_query has the other one, embed_query(question) -> vector. Those two callables are the only difference between this episode and the self-hosted ones.

What the pipeline asks the model

run_single_qa in wattbot_rag.py makes two calls per question. The first asks for a JSON object in the WattBot submission format, from the retrieved context only; the second, made only when the first produced an answer, asks for a one-to-three-sentence explanation of how the cited material supports it. The system prompt of the first call is the contract that makes the output machine-readable:

PYTHON 

print(ANSWER_SYSTEM_PROMPT)

A question is marked unanswerable (is_blank everywhere, as the competition requires) when the model says so or when the best retrieval score is below retrieval_threshold.

Step 11 – Run the WattBot evaluation with Bedrock

Now we can loop over the questions, run retrieval + Bedrock generation, and write a wattbot_solutions_bedrock.csv file (locally and to your bucket). train_QA.csv has 245 questions, and at two calls each that is about twenty minutes. For class we score a stratified sample instead. sample_questions_by_type draws ten questions from each evidence type (Quote, Table, Figure, Math, unanswerable, cross-paper, reconcile), 53 questions after overlaps and the two missing papers, so every row of the breakdown in the next step has something in it. The first 53 rows of the file would be nearly all Quote questions. The seed fixes the draw, and the Processing Jobs episode uses the same one, so the two runs are compared on the same questions. Set per_type=0 to run the whole file.

The Processing Jobs episode reuses this same loop with self-hosted models; the only difference there is that the embedding, answer and explanation phases call local models instead of Bedrock-hosted ones.

PYTHON 

from wattbot_rag import sample_questions_by_type

eval_df = sample_questions_by_type(train_df, per_type=10, seed=0)   # per_type=0 for all 220 remaining questions
print(f"Scoring {len(eval_df)} of {len(train_df)} questions")

results = []
run_start = time.time()

for _, row in eval_df.iterrows():
    print("#" * 96)
    print(f"QUESTION: {row['question']}")

    out = run_single_qa(
        row=row,                             # one train_QA.csv row: id, question, expected answer_unit
        embed_query=embed_query,             # the two model callables: the only episode-specific inputs
        generate=call_bedrock_chat,
        chunk_embeddings=chunk_embeddings,   # the index: the matrix plus the chunks it was built from
        chunked_docs=chunked_docs,
        docid_to_url=docid_to_url,           # for the ref_url column
        top_k=8,                             # chunks handed to the model as context
        retrieval_threshold=0.25,            # best cosine score below this => unanswerable (is_blank)
        verbose=True,
    )

    print(f"ANSWER: {out['answer']}")
    print(f"ref_ids: {out['ref_id']}")
    print(f"EXPLANATION: {out['explanation']}")
    results.append(out)

results_df = pd.DataFrame(results)
print(f"Answered {len(results_df)} questions in {(time.time() - run_start) / 60:.1f} min")

output_dir = "outputs"
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, "wattbot_solutions_bedrock.csv")
results_df.to_csv(output_path, index=False)
s3.upload_file(output_path, bucket_name, solutions_key)
print(f"Wrote predictions to {output_path} and s3://{bucket_name}/{solutions_key}")

results_df.head()

Step 12 – Score against the training answers

WattBot scores each prediction with three components, 0.75 for answer_value (numeric within ±0.1% relative tolerance, categorical match after normalization, and zero if you answer a question the corpus cannot answer), 0.20 for the F1 between your ref_id set and the ground-truth citations, and 0.05 for refusing cleanly, with every evidence field blank, when the answer is is_blank. The number comes from Score.py, the competition’s own metric file, which compute_wattbot_score in wattbot_rag.py calls for the headline score so that what you see here is what the leaderboard would report. It also prints the share of correct answers by evidence type, the seven flag columns of train_QA.csv, which are Quote (the answer is a sentence in a paper), Table, Figure, Math (a calculation from numbers in the text), is_NA (the corpus cannot answer it, and the right response is is_blank), CrossPaper (numbers from two papers have to be combined) and Reconcile (two figures that sound alike but measure different things, such as water withdrawn against water consumed, or a popular claim against a measured value, and the answer is the ratio or the distinction between them; quoting either figure alone is the wrong answer the question is built to catch). A question can carry several. Read that table before changing anything. Figure and table questions failing while quote questions pass points at text extraction, which loses figures and flattens tables, not at the model. The wrapper adds a per-row score, joins on id, and prints the worst-scoring rows so you can see why the pipeline lost points.

PYTHON 

scored_df = compute_wattbot_score(
    train_qa_path=train_qa_path,    # ground truth
    preds_path=output_path,         # our predictions, joined to it on id
    n_examples=20,
)

Why were the wrong ones wrong?

The score tells you which questions failed, not why. Two different failures look the same in that table. Either retrieval never surfaced the right paper, so the model had nothing to work with, or the right passage was in front of it and it still declined or misread it. The cell below separates them for the lowest-scoring rows. It joins the model’s own explanation and the passages it cited back onto the score, then reruns retrieval for the question (one embedding call) and reports whether the ground-truth paper was among the top-k chunks and what the best cosine score was.

PYTHON 


worst = scored_df.sort_values("wattbot_score").head(8)
by_id = results_df.set_index("id")   # the model's answer, explanation and cited passages, keyed by question id

for _, row in worst.iterrows():
    qid = row["id"]
    gt_refs = set(split_ref_ids(row["ref_id_gt"]))   # train_QA.csv stores the list as a string
    retrieved = retrieve_top_k(embed_query(row["question_gt"]), chunk_embeddings, chunked_docs, k=8)   # the merge suffixed both sides' columns
    retrieved_docs = [r["doc_id"] for r in retrieved]
    print("=" * 96)
    print(f"{qid}: {row['question_gt']}")
    print(f"  expected: {row['answer_value_gt']}  from {sorted(gt_refs)}")
    print(f"  predicted: {row['answer_value_pred']}  citing {row['ref_id_pred']}")
    print(f"  retrieval: top cosine {retrieved[0]['score']:.3f}; top-8 papers {retrieved_docs}")
    print(f"  ground-truth paper retrieved: {'yes' if gt_refs & set(retrieved_docs) else 'NO'}")
    if by_id.loc[qid, "explanation"] != "is_blank":
        print(f"  model's explanation: {by_id.loc[qid, 'explanation']}")
        print(f"  passages it cited: {str(by_id.loc[qid, 'supporting_materials'])[:400]}")

Read the two lines at the bottom of each block together. “Ground-truth paper retrieved: NO” is a retrieval problem, where a different chunk size, a reranker, or a better embedding model is the fix, and no generator will help. “Yes” with is_blank is the model declining on evidence it had, which the system prompt and the retrieval threshold control. “Yes” with a wrong value is a reading error, and the cited passage usually shows what it latched onto, a nearby number, a different unit, a different model in the same table. And sometimes the explanation shows the model read correctly and the label is the problem, as with a question asking what inference “can account for” labeled 90 where three papers in the corpus give 60, 70 and 80 to 90 percent, so a range is arguably the better answer, or a fact that appears in several papers with the label naming only one of them. Labels are written by people. When the explanation convinces you, the fix is to the label, and the competition organizers want to hear about it. Each of those is a different lever, and the score alone would have you pulling the wrong one.

Challenge

What did that run cost?

Every converse response carried a usage block, and call_bedrock_chat kept them all in bedrock_usage. Sum the input and output tokens over the evaluation run and price them at the current on-demand rates for DeepSeek V3 on the Bedrock pricing page (fill the two prices in below). Add the embedding step, ~9,000 chunks of 1,200 characters at Titan Text Embeddings V2’s $0.02 per million input tokens. How does the total compare with the overview’s estimate for “the whole WattBot pipeline on Bedrock”, and with the ml.g5.xlarge GPU the Processing Jobs episode rents for the same questions, $1.41 an hour?

The two prices are from the Bedrock pricing page for DeepSeek on-demand in October 2026; update them for the model and region you used. The rest sums the usage blocks and prints the cost.

PYTHON 

price_in_per_m = 0.62    # $ per million input tokens: DeepSeek V3 on-demand, October 2026; check the pricing page for your model and region
price_out_per_m = 1.85   # $ per million output tokens, same source

in_tokens = sum(u["inputTokens"] for u in bedrock_usage)
out_tokens = sum(u["outputTokens"] for u in bedrock_usage)
gen_cost = in_tokens / 1e6 * price_in_per_m + out_tokens / 1e6 * price_out_per_m

# Titan bills per input token; ~4 characters per token is a fair estimate for English prose
embed_tokens = sum(len(ch["text"]) for ch in chunked_docs) / 4
embed_cost = embed_tokens / 1e6 * 0.02

print(f"{len(bedrock_usage)} generation calls: {in_tokens:,} input + {out_tokens:,} output tokens -> ${gen_cost:.3f}")
print(f"~{embed_tokens:,.0f} embedding tokens -> ${embed_cost:.3f}")
print(f"Total ~${gen_cost + embed_cost:.2f}; the Processing Jobs episode rents an ml.g5.xlarge at $1.41 an hour for the same questions")

With 53 questions, two calls each and ~1,500-token prompts, the run was 140K input and 9K output tokens, ten cents at DeepSeek V3’s prices, plus five cents of embeddings for the 2.4 million tokens of corpus. All 245 questions is about half a dollar, which is the overview’s estimate for the whole pipeline. The Processing Jobs episode answers the same 53 questions on a rented ml.g5.xlarge and bills about $0.50 of GPU time for them, three times this run, most of it loading a model that Bedrock already has loaded.

Challenge

Swap the models

Everything above is two model IDs and two callables. The model catalog has dozens of chat models and several embedding models, and “Choosing the two models” at the top of the episode said the leaderboards cannot tell you which one wins on your questions. Find out.

  1. A different generator. Pick another chat model from the catalog (for example us.amazon.nova-2-lite-v1:0, us.meta.llama4-maverick-17b-instruct-v1:0, or a Mistral or Claude model if your account has access), wrap it in a tagged profile as in Step 7, and rerun the Step 11 loop on the same 53-question sample with call_bedrock_chat pointed at it. Score it with compute_wattbot_score and price it from its usage blocks. Which evidence types moved, and what did a point of score cost?
  2. A different embedder. Retrieval is the other half. The Ohio catalog also lists Cohere Embed v4 (cohere.embed-v4:0), and its request body is different from Titan’s, a list of texts and an input_type that is search_document for the corpus and search_query for the question. Embed the chunks with it under a new file name so the Titan matrix is kept (ten minutes at the quota, so not in class), rerun Step 11 with DeepSeek unchanged, and score.

Keep the rest fixed. One change per experiment, or you will not know what moved the score.

Both solution cells start with %%script false --no-raise-error, which makes Jupyter skip the cell, so a Run All of the episode does not spend ten minutes here; delete that line to run one.

Generator. call_bedrock_chat already takes a model_id, so the swap is a profile and a lambda:

PYTHON 

%%script false --no-raise-error
# Challenge solution: delete the line above to run this cell. It is skipped on Run All so the episode's own cells are not delayed by it.
alt_generation_model_id = "us.amazon.nova-2-lite-v1:0"   # any chat model in the catalog with a converse API
alt_model_arn = tagged_profile_arn(bedrock, alt_generation_model_id, "wattbot-gen-alt", name, project)

usage_before = len(bedrock_usage)
alt_results = []
for _, row in eval_df.iterrows():
    alt_results.append(run_single_qa(
        row=row, embed_query=embed_query,
        generate=lambda s, u, max_new_tokens=512: call_bedrock_chat(s, u, max_new_tokens, model_id=alt_model_arn),
        chunk_embeddings=chunk_embeddings, chunked_docs=chunked_docs, docid_to_url=docid_to_url,
        top_k=8, retrieval_threshold=0.25,
    ))
alt_path = os.path.join(output_dir, "wattbot_solutions_bedrock_alt.csv")
pd.DataFrame(alt_results).to_csv(alt_path, index=False)

alt_usage = bedrock_usage[usage_before:]   # only this run's calls
print(f"{sum(u['inputTokens'] for u in alt_usage):,} input + {sum(u['outputTokens'] for u in alt_usage):,} output tokens; price them from the pricing page")
compute_wattbot_score(train_qa_path=train_qa_path, preds_path=alt_path, n_examples=5)

Read the per-type lines against the DeepSeek run. A cheaper model that holds its own on Quote questions and loses on Math and CrossPaper is the usual pattern; whether that trade is worth it depends on which questions you have.

Embedder. First list the embedding models the region’s catalog has, because model IDs come and go and an ID that is not in the list fails at the profile step with “The provided model identifier is invalid”:

PYTHON 

for m in bedrock.list_foundation_models(byOutputModality="EMBEDDING")["modelSummaries"]:
    print(f"{m['modelId']:<45} {m['providerName']}")

Cohere’s request takes a list of texts and says whether they are documents or a query; its response is a list of vectors, nested under the number type asked for. One wrinkle. The catalog lists cohere.embed-v4:0 for Ohio, but creating a profile on it fails with “does not support On Demand inference”, because in this region the model is served only through the US cross-region profile, us.cohere.embed-v4:0. tagged_profile_arn catches that error and switches to the us. ID itself, which is the same thing that happens for Llama 4 and several other models. Everything else is Step 8 again, a tagged profile, one call per chunk, a cached matrix, and a query function that uses the same model.

PYTHON 

%%script false --no-raise-error
# Challenge solution: delete the line above to run this cell. Ten minutes of embedding, so not in class.
cohere_model_id = "cohere.embed-v4:0"   # listed for Ohio but served only through the US cross-region profile; tagged_profile_arn switches to us.cohere.embed-v4:0 on its own
cohere_arn = tagged_profile_arn(bedrock, cohere_model_id, "wattbot-embed-cohere", name, project)

def cohere_embed(text: str, input_type: str = "search_document") -> np.ndarray:
    """One Cohere Embed v4 vector; input_type is search_document for the corpus, search_query for a question."""
    body = {"texts": [text[:2000]], "input_type": input_type, "embedding_types": ["float"]}
    response = bedrock_runtime.invoke_model(modelId=cohere_arn, contentType="application/json", accept="application/json", body=json.dumps(body))
    return np.asarray(json.loads(response["body"].read())["embeddings"]["float"][0], dtype="float32")

cohere_emb_path = data_dir / "embeddings_cohere.npy"   # a new file: the Titan matrix stays
if cohere_emb_path.exists():
    cohere_embeddings = np.load(cohere_emb_path)
else:
    t0 = time.time()
    cohere_embeddings = np.vstack([cohere_embed(ch["text"]) for ch in chunked_docs])
    np.save(cohere_emb_path, cohere_embeddings)
    print(f"Embedded {len(chunked_docs)} chunks in {(time.time() - t0) / 60:.1f} min")

def embed_query_cohere(question: str) -> np.ndarray:
    return cohere_embed(question, input_type="search_query")

Then the Step 11 loop with embed_query=embed_query_cohere and chunk_embeddings=cohere_embeddings, DeepSeek unchanged, and the scorer. The hit-rate cell from Step 13 is the sharper instrument here. Compute it for the new matrix before scoring, because an embedder that finds the paper less often cannot be rescued by any generator. A cheaper variant of the same experiment keeps Titan and adds "dimensions": 256 to its request body in Step 8, which asks whether a four-times smaller index retrieves as well.

Step 13 – Multimodal embeddings for the table and figure questions

Look back at the evidence-type table in Step 12. The Table and Figure rows are usually the worst, and the reason is upstream of the model. pypdf gives us a page’s text, so a table comes out as a run of numbers with its column headers somewhere else on the page, and a figure comes out as its caption and nothing more. A question whose answer is a bar in a plot has nothing in the chunks to find, and a question whose answer is a cell in a table is being read from a flattened copy.

A multimodal embedding model is one model that maps text and images into the same vector space, so a text question can be compared with an image directly, and the picture of a table lands near a sentence about its contents. Bedrock has two from Amazon, Amazon Nova Multimodal Embeddings (text, images, documents, video and audio, released 2025) and the older Titan Multimodal Embeddings G1 (text and images). Both are per-call, like Titan Text, and billed per image plus per text token; the Bedrock pricing page has the rates, on the order of a hundredth of a cent per image. One catch. In October 2026 Nova Multimodal Embeddings is offered in us-east-1 and not in Ohio (check the model catalog for that region). A Bedrock call is an HTTPS request and the region is a parameter, so from Ohio that is a second client pointed at Virginia; it also means the images leave the region, which is the thing to check against your institution’s data policy before doing this with documents that are not public papers.

This step cuts the tables and figures out of the PDFs as images, builds a second index from them next to the text one, measures whether it finds the right paper for table and figure questions more often, and then hands the images to a model that can look at them. Each cell is short because the plumbing is in wattbot_rag.py.

Extract the tables and figures

Why cut them out at all, instead of handing the embedding model a whole page, or the PDF itself (Nova accepts documents as input)? For the same reason Step 6 cut the text into chunks instead of embedding each paper whole. An embedding is a fixed-length vector, a thousand numbers, whatever goes in; the more that goes in, the less of each thing survives in it. A page holding two tables, a figure and three paragraphs gets one vector that is a blur of all six, and a question about one of the tables has to match that blur. A crop of the table gets a vector that is about the table. The same holds at generation time. The vision model reads a crop at full resolution for a few hundred tokens, where a page costs a thousand tokens and shows the table at a quarter of the size. We tried whole pages first on this corpus. They helped figure questions a little and tables not at all. The crops, as the callout at the end shows, help both.

A PDF has no table or figure objects to pull out. It stores text runs, raster images, and drawing commands (lines, fills, curves), and a plot is a few thousand of the last. PyMuPDF ships three detectors that turn those back into boxes. find_tables finds tables from their rulings or their text alignment, get_image_info lists the embedded raster images, and cluster_drawings groups the drawing commands of a vector figure into one bounding box. extract_figures_and_tables in wattbot_rag.py runs the three on every page, drops page-sized background fills, rules and logos, pairs each box with the nearest caption line starting “Table N” or “Figure N” (tables are captioned above, figures below) and widens the crop to include that caption, so the image carries the identifier a question will use. A caption whose box no detector found falls back to the whole page, so nothing the text mentions is missed. Open the function; it is about eighty lines, and every rule in it is a line you can change. The detectors are the slow part, a second or two per page on this instance, so the cell runs them only on the pages whose text carries a caption, under half of the corpus; in class, with use_instructor_embeddings set, it downloads the crops your instructor made instead.

PYTHON 

%pip install -q pymupdf

PYTHON 

from wattbot_rag import extract_figures_and_tables, find_caption_pages, embed_multimodal, image_blocks, hit_at_k

crop_dir = data_dir / "crops"
crops_path = data_dir / "crops.jsonl"
shared_crops_key = "wattbot/crops.zip"   # the instructor's crops, staged like the embeddings

if not crops_path.exists() and use_instructor_embeddings and _s3_object_exists(shared_bucket, shared_crops_key):
    s3.download_file(shared_bucket, shared_crops_key, str(data_dir / "crops.zip"))
    with zipfile.ZipFile(data_dir / "crops.zip") as zf:
        zf.extractall(data_dir)
    (data_dir / "crops.zip").unlink()
    print("Downloaded the instructor's crops")

if crops_path.exists():
    crops = [json.loads(line) for line in open(crops_path)]
else:
    caption_pages: Dict[str, Dict[int, set]] = {}   # per paper, the pages whose text carries a caption, and which kinds: the detectors run only there
    for p in find_caption_pages(page_docs):
        caption_pages.setdefault(p["doc_id"], {})[p["page_num"]] = {"table" if c.startswith("Table") else "figure" for c in p["captions"]}
    t0 = time.time()
    crops = []
    for i, (doc_id, page_nums) in enumerate(sorted(caption_pages.items())):
        crops += extract_figures_and_tables(os.path.join(PDF_DIR, f"{doc_id}.pdf"), doc_id, crop_dir, pages=page_nums)
        if (i + 1) % 10 == 0:
            print(f"{i + 1} papers, {len(crops)} crops so far ({time.time() - t0:.0f} s)")
    with open(crops_path, "w") as f:
        f.writelines(json.dumps(c) + "\n" for c in crops)
    print(f"Extracted {len(crops)} crops from {len(caption_pages)} papers in {(time.time() - t0) / 60:.1f} min")

n_tables = sum(c["kind"] == "table" for c in crops)
n_pages = sum(c["whole_page"] for c in crops)
print(f"{n_tables} tables and {len(crops) - n_tables} figures; {n_pages} of them are whole-page fallbacks")
crops[:3]

Look at a few before trusting them. A crop that is a logo, a page rule, or half a plot tells you which threshold to change.

PYTHON 

from IPython.display import Image, display

for c in [c for c in crops if c["kind"] == "table"][:2] + [c for c in crops if c["kind"] == "figure"][:2]:
    print(f"{c['doc_id']} page {c['page_label']}: {c['caption']}")
    display(Image(filename=c["png_path"], width=500))

Embed the crops

The embedding model lives in us-east-1, so this is a second runtime client pointed there. One difference from the text models is that Bedrock rejects this model’s request body when it arrives through an application inference profile, so these calls use the model ID directly and are the one Bedrock spend in this lesson that carries no project tags. embed_multimodal builds the model’s request body (Nova’s says whether a vector is for the index or for a query, and that the image is a document, not a photo) and returns a unit vector. Set multimodal_model_id to amazon.titan-embed-image-v1 to try the older model; the function handles both.

PYTHON 

multimodal_region = "us-east-1"
multimodal_model_id = "amazon.nova-2-multimodal-embeddings-v1:0"   # called by ID: this model does not accept calls through an inference profile

bedrock_runtime_mm = boto3.client("bedrock-runtime", region_name=multimodal_region,
                                  config=Config(retries={"max_attempts": 10, "mode": "adaptive"}))

v = embed_multimodal(bedrock_runtime_mm, multimodal_model_id, image_png=Path(crops[0]["png_path"]).read_bytes())
print(f"{crops[0]['doc_id']} {crops[0]['caption'][:40]!r} -> vector of {v.shape[0]} dimensions")

One call per crop, cached like the text matrix, locally, then in your bucket, then computed. The loop runs six calls at a time through a thread pool and prints every fifty; the 2,213 crops take about four minutes. The question side uses the same model with purpose="query".

PYTHON 

crop_emb_path = data_dir / "embeddings_crops_multimodal.npy"
crop_emb_key = "wattbot/embeddings/embeddings_crops_multimodal.npy"

if crop_emb_path.exists():
    crop_embeddings = np.load(crop_emb_path)
elif use_instructor_embeddings and _s3_object_exists(shared_bucket, "wattbot/embeddings_crops_multimodal.npy"):
    s3.download_file(shared_bucket, "wattbot/embeddings_crops_multimodal.npy", str(crop_emb_path))
    crop_embeddings = np.load(crop_emb_path)
    print("Downloaded the instructor's crop embeddings")
elif _s3_object_exists(bucket_name, crop_emb_key):
    s3.download_file(bucket_name, crop_emb_key, str(crop_emb_path))
    crop_embeddings = np.load(crop_emb_path)
else:
    from concurrent.futures import ThreadPoolExecutor

    def embed_crop(c):
        return embed_multimodal(bedrock_runtime_mm, multimodal_model_id, image_png=Path(c["png_path"]).read_bytes())

    t0 = time.time()
    rows = []
    with ThreadPoolExecutor(max_workers=6) as pool:   # one image per call: six in flight at a time
        for i, vec in enumerate(pool.map(embed_crop, crops)):
            rows.append(vec)
            if (i + 1) % 50 == 0:
                print(f"Embedded {i + 1} / {len(crops)} crops ({time.time() - t0:.0f} s, {(i + 1) / (time.time() - t0):.1f} per second)")
    crop_embeddings = np.vstack(rows)
    np.save(crop_emb_path, crop_embeddings)
    s3.upload_file(str(crop_emb_path), bucket_name, crop_emb_key)
    print(f"Embedded {len(crops)} crops in {(time.time() - t0) / 60:.1f} min")

assert crop_embeddings.shape[0] == len(crops), "crop embeddings and crops are out of step; delete the cached file and rerun"
print("Crop embeddings shape:", crop_embeddings.shape)

def embed_query_mm(question: str) -> np.ndarray:
    return embed_multimodal(bedrock_runtime_mm, multimodal_model_id, text=question, purpose="query")

Does the image index find the right table or figure?

Retrieval is the same retrieve_top_k as Step 9, over the crop matrix instead of the chunk matrix. Two measures, both hit rates, higher is better, 1.00 meaning every question had it. The first is coarse. Is a ground-truth paper among the top results? train_QA.csv names the paper for every question. The second is the one that matters for this step. Is the table or figure the answer sits in among the crops the model would be shown? Most table and figure questions name it in supporting_materials (“Table 3”, “Fig. 5, upper panel”), and cited_captions pulls that identifier out; crop_matches checks a crop’s caption against it. For the text index, the crops it “shows” are those on the pages its top chunks came from, which is what the third vision run below attaches. Quote-only questions are the control for the paper measure, where the image index should have no advantage.

PYTHON 

from wattbot_rag import cited_captions, crop_matches

k = 8
groups = {"Table": train_df[train_df["Table"] == 1], "Figure": train_df[train_df["Figure"] == 1],
          "Quote only": train_df[(train_df["Quote"] == 1) & (train_df["Table"] == 0) & (train_df["Figure"] == 0)].head(40)}
crops_by_page: Dict[tuple, list] = {}
for c in crops:
    crops_by_page.setdefault((c["doc_id"], c["page_num"]), []).append(c)

def crops_behind(chunks):   # the crops on the pages the retrieved chunks came from
    return [c for ch in chunks for c in crops_by_page.get((ch["doc_id"], ch["page_num"]), [])]

print("Hit rate (higher is better). Paper: a ground-truth paper is in the top 8 results. Table/figure: the one the answer")
print("sits in is among the crops shown (the image index's top 3; for the text index, the crops on its top chunks' pages)\n")
print(f"{'questions':<12}{'n':>4}{'paper, text':>14}{'paper, image':>14}{'crop, text':>14}{'crop, image':>14}")
for label, qs in groups.items():
    gt = [split_ref_ids(r) for r in qs["ref_id"]]
    text_chunks = [retrieve_top_k(embed_query(q), chunk_embeddings, chunked_docs, k=k) for q in qs["question"]]
    image_crops = [retrieve_top_k(embed_query_mm(q), crop_embeddings, crops, k=3) for q in qs["question"]]
    paper_text = hit_at_k([[c["doc_id"] for c in r] for r in text_chunks], gt)
    paper_image = hit_at_k([[c["doc_id"] for c in r] for r in image_crops], gt)
    # the crop measure, over the questions that name their table or figure
    targets = [(refs, cited_captions(sm)) for refs, sm in zip(gt, qs["supporting_materials"])]
    crop_text, crop_image, n_named = [], [], 0
    for (refs, labels), chunks_i, crops_i in zip(targets, text_chunks, image_crops):
        if not labels:
            continue
        n_named += 1
        crop_text.append(any(crop_matches(c, d, l) for c in crops_behind(chunks_i) for d in refs for l in labels))
        crop_image.append(any(crop_matches(c, d, l) for c in crops_i for d in refs for l in labels))
    crop_cols = f"{np.mean(crop_text):>14.2f}{np.mean(crop_image):>14.2f}  (of {n_named} that name one)" if n_named else f"{'':>14}{'':>14}"
    print(f"{label:<12}{len(qs):>4}{paper_text:>14.2f}{paper_image:>14.2f}{crop_cols}")

Read the paper columns first. Both indexes find the paper for nearly every table and figure question, so at that level there is nothing to win, and the image index is behind. The crop columns are the real comparison. They say whether the text route, chunks and the crops on their pages, or the image route, the crops themselves, more often puts the exact table or figure in front of the model. If the image column is higher there while the paper column is lower, the image index is doing a different job than the text index, finding the object, not the document, and the next cells show whether the model can use it. If both crop columns are low, the fix is upstream, in the extraction, where a table the detector missed became a whole-page fallback, and a page is a weak key.

Let the model see the table or figure

Finding the table is half of it; DeepSeek V3 reads text only, so a figure retrieved as an image is no use to it. converse accepts image blocks in a user message for models that can see, and the swap is the model ID. Amazon Nova 2 Lite reads images and documents and is in the catalog for Ohio through the cross-region profile below (check the model catalog; us.amazon.nova-lite-v1:0 is the cheaper, weaker alternative). call_vision_chat is call_bedrock_chat with an images argument. For each attached crop it adds a text label naming the paper and the caption, then the image, so the model can cite what it reads.

PYTHON 

base_vision_model_id = "us.amazon.nova-2-lite-v1:0"
vision_model_id = tagged_profile_arn(bedrock, base_vision_model_id, "wattbot-vision", name, project)
vision_usage = []

def call_vision_chat(system_prompt: str, user_prompt: str, max_new_tokens: int = 512, images=()) -> str:
    """One system + user turn, with zero or more (label, png_bytes) images attached to the user turn."""
    content = [{"text": user_prompt}]
    if images:
        content.append({"text": "The images below are tables and figures from the same papers; read values from them where the text above is incomplete."})
        content.extend(image_blocks(images))
    response = bedrock_runtime.converse(
        modelId=vision_model_id,
        system=[{"text": system_prompt}],
        messages=[{"role": "user", "content": content}],
        inferenceConfig={"maxTokens": max_new_tokens, "temperature": 0},
    )
    vision_usage.append(response["usage"])
    return "\n".join(b["text"] for b in response["output"]["message"]["content"] if "text" in b).strip()

def crop_label(c) -> str:
    return f"{c['doc_id']}, page {c['page_label']}, {c['caption'][:60]}" if c["caption"] else f"{c['doc_id']}, page {c['page_label']}"

print(call_vision_chat("You are a concise assistant.", "What is in the attached image? One sentence.", 64,
                       images=[(crop_label(crops[0]), Path(crops[0]["png_path"]).read_bytes())]))

Three runs over the Table and Figure questions, all through the vision model, so that the only thing changing between them is what the model is shown. The first is the control, text chunks only. The second attaches the top three crops from the image index. The third attaches the crops from the pages the top text chunks came from. If the hit-rate table says the text index is the better finder, the question is whether finding with text and reading with the image is the better split. Each run is the same loop as Step 11, so every question costs an answer call and an explanation call; our runs took 1.6, 4.0 and 3.0 minutes, the image runs longer because every attached crop is hundreds of input tokens. In class, with use_instructor_embeddings set, the cell downloads the three result files from your instructor’s run of it instead and the next cell scores them; running the loops yourself is for after class.

PYTHON 

tf_df = train_df[(train_df["Table"] == 1) | (train_df["Figure"] == 1)]
print(f"{len(tf_df)} table and figure questions; each run makes an answer call and an explanation call per question, so about {6 * len(tf_df)} calls for three runs")

def crops_for(question: str, images_from: str):
    """The crops to attach: none, the image index's top three, or those on the pages behind the top text chunks."""
    if images_from == "image_index":
        return retrieve_top_k(embed_query_mm(question), crop_embeddings, crops, k=3)
    if images_from == "text_index":
        found = []
        for c in crops_behind(retrieve_top_k(embed_query(question), chunk_embeddings, chunked_docs, k=8)):
            if c not in found:
                found.append(c)
        return found[:3]
    return []

def run_tf(images_from: str) -> pd.DataFrame:
    rows = []
    for i, (_, row) in enumerate(tf_df.iterrows(), start=1):
        images = [(crop_label(c), Path(c["png_path"]).read_bytes()) for c in crops_for(row["question"], images_from)]
        generate = lambda s, u, max_new_tokens=512, images=images: call_vision_chat(s, u, max_new_tokens, images=images)
        rows.append(run_single_qa(row=row, embed_query=embed_query, generate=generate, chunk_embeddings=chunk_embeddings,
                                  chunked_docs=chunked_docs, docid_to_url=docid_to_url, top_k=8, retrieval_threshold=0.25))
        if i % 10 == 0:
            print(f"  {i} / {len(tf_df)} questions ({time.time() - t0:.0f} s)")
    return pd.DataFrame(rows)

vision_runs = {"text chunks only": "none", "chunks + crops from the image index": "image_index", "chunks + crops behind the text chunks": "text_index"}
vision_files = {images_from: f"wattbot_solutions_vision_{images_from}.csv" for images_from in vision_runs.values()}

t0 = time.time()
for label, images_from in vision_runs.items():
    local_path = os.path.join(output_dir, vision_files[images_from])
    shared_key = f"wattbot/solutions/{vision_files[images_from]}"
    if use_instructor_embeddings and _s3_object_exists(shared_bucket, shared_key):   # in class: the instructor's run of this cell, scored in the next
        s3.download_file(shared_bucket, shared_key, local_path)
        print(f"{label}: downloaded the instructor's results")
        continue
    print(f"{label}: running")
    run_tf(images_from).to_csv(local_path, index=False)
    s3.upload_file(local_path, bucket_name, shared_key)   # kept in your bucket too, next to the Step 11 predictions
    print(f"{label}: done at {(time.time() - t0) / 60:.1f} min")

Score all three the same way as Step 12. The scorer compares the rows it finds, so the breakdown is over these questions only; the Table and Figure lines are the ones to read, across the three runs and against the DeepSeek run.

PYTHON 

for label, images_from in vision_runs.items():
    print("=" * 96)
    print(f"Nova 2 Lite, {label}")
    compute_wattbot_score(train_qa_path=train_qa_path, preds_path=os.path.join(output_dir, f"wattbot_solutions_vision_{images_from}.csv"), n_examples=3)
Callout

What to expect, and what it cost

Measured once on this corpus, 50 table and figure questions, Nova 2 Lite throughout. At the paper level the text index wins. A ground-truth paper is in its top 8 for 94% of table questions and 94% of figure questions, against 77% and 71% for the image index, and for 82% against 55% of quote-only questions, where a picture has nothing to offer. At the crop level the two routes are even. The crop the answer sits in is among those shown for 55% of table questions by the text route and 52% by the image index, and for 57% against 50% of figure questions. Then reading. Text chunks alone scored 0.64, with 57% of table answers and 59% of figure answers correct. Chunks plus the image index’s top three crops scored 0.75, with 69% of tables and 71% of figures correct and citation F1 up from 0.85 to 0.93. Chunks plus the crops on the pages behind the top text chunks, the route the paper-level hit rate seemed to favor, scored 0.67. Two readings. Showing the model the table or figure is worth about ten points on the questions that depend on one, tables included, which whole pages at 100 dpi had not managed. And the crop-level hit rates say where the image index’s lead does not come from. It finds the answer’s own crop no more often than the text route does. The difference is in what else the model is shown, three crops chosen by their similarity to the question against every crop on the pages behind the top chunks. Fifty questions is a small sample, so read the direction, not the decimals. For the bill, every crop is a few hundred input tokens at Nova 2 Lite’s rate on the pricing page, so a run with images costs two to three times the control and takes three to four times as long, 1.6, 4.0 and 3.0 minutes for the three runs here. vision_usage holds all three runs’ usage blocks, and the embedding of the 2,213 crops is the per-image rate times that count, a few cents.

Wrap-up

You now have a complete WattBot evaluation on Bedrock, billed per token, with nothing left running. The Processing-jobs episode reruns it on a self-hosted model, and the optional GPU-notebook episode once more by hand, so by the end there are three versions of the same pipeline to compare, Bedrock, a Processing job and a notebook GPU. The code they share is the retrieval, prompting and scoring in wattbot_rag.py; what each one implements is the two callables, embed_query and generate.

Use Bedrock or another hosted API when a catalog model will do, when the number of questions is modest or you are still prototyping, and when you would rather pay per token than manage a GPU. Use a self-hosted model on a GPU when the work needs a model the catalog does not carry, when you will run large batches repeatedly, when you need control over the exact checkpoint, or when you already have GPUs you are not paying for by the hour.

The model is one part of a RAG system, and not the part where most of the accuracy is won. A bigger or newer model does not fix weak retrieval. If the chunks do not match the questions, a strong model still has nothing to read. The gains over time come from the work around the model, chunking, metadata and filtering, reranking, post-processing, and the score breakdown by evidence type in Step 12 says which of them to work on. A bigger model or a longer prompt raises the bill per question as well as the score, and Bedrock lets you measure that trade by changing one model ID.

Key Points
  • A complete RAG pipeline on AWS needs no GPU. A CPU notebook plus Bedrock-hosted embedding and generation models is enough.
  • Corpus preparation (unpacking the staged PDFs, page extraction, chunking) is cheap CPU work; cache chunks.jsonl on S3 so later episodes and runs can reuse it.
  • Bedrock calls carry no tags; route them through a tagged application inference profile or the spend is anonymous on the bill.
  • converse is the provider-neutral way to call a chat model on Bedrock; its usage block is what you are billed on. DeepSeek V3 is the generator here, and any catalog model is a one-line swap.
  • Bedrock bills per token and leaves nothing running between calls, which is why it is the default route for RAG inference.
  • The model is one tunable component; most accuracy gains come from chunking, retrieval, and post-processing.