Every client conversation about generative AI arrives at the same place within twenty minutes: "we don't want the model to make things up about our documents." That is a retrieval problem before it is a language-model problem. Retrieval-augmented generation (RAG) answers a question by first finding the handful of passages in your corpus that are actually relevant, then asking a language model to answer using only those passages, with citations.
This tutorial builds a working RAG pipeline entirely on the TensorFlow/Keras stack: chunking, embeddings with KerasHub, an approximate-nearest-neighbour index with ScaNN, grounded generation with Gemma, and - the part most tutorials skip - an evaluation harness that tells you whether any of it works.
It complements two other tutorials here. Fine-tuning Gemma with KerasHub changes what the model knows how to do; RAG changes what the model knows about, without retraining. And Two-Tower Retrieval uses the same ScaNN machinery for recommenders, where you train the encoders; here we start with a pretrained text encoder and no labelled data at all.
Fine-tune or retrieve?
Decide this before you write code, because picking wrong costs weeks:
| Symptom | Reach for |
|---|---|
| Model doesn't know your facts, policies, prices, tickets | RAG |
| Facts change weekly | RAG (reindex, don't retrain) |
| You need citations and auditability | RAG |
| Model knows the facts but answers in the wrong style/format | Fine-tune |
| A specialised task: classification, extraction, structured output | Fine-tune |
| Both | RAG first, fine-tune later if the gap persists |
RAG is cheaper, faster to ship, and easier to defend to a compliance team, because you can point at the source of every sentence. Start there.
Setup
pip install -U keras keras-hub scann tensorflow
import os
os.environ["KERAS_BACKEND"] = "tensorflow"
import keras, keras_hub, numpy as np, tensorflow as tf
1. Chunking: the step that decides your ceiling
Retrieval quality is capped by chunk quality. A chunk should be big enough to answer a question on its own and small enough that it is about one thing.
import re
def chunk_document(text, doc_id, target_chars=1200, overlap_chars=200):
"""Split on paragraph boundaries, pack up to target size, overlap a little."""
paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
chunks, buf = [], ""
for para in paragraphs:
if len(buf) + len(para) > target_chars and buf:
chunks.append(buf)
buf = buf[-overlap_chars:] + "\n\n" + para
else:
buf = f"{buf}\n\n{para}" if buf else para
if buf:
chunks.append(buf)
return [
{"doc_id": doc_id, "chunk_id": f"{doc_id}#{i}", "text": c}
for i, c in enumerate(chunks)
]
Field notes:
- Respect structure. Split Markdown or HTML on headings, not at a fixed character count. A chunk that starts mid-sentence retrieves badly and reads worse in a citation.
- Prepend context. Store
"{document title} > {section heading}\n\n{chunk text}"as the embedded text. It is a one-line change and it measurably improves retrieval on corpora with repetitive boilerplate. - Keep tables and code blocks whole. Half a table is worse than no table.
- Carry metadata: source URL, section, last-modified date, access level. You will need every one of these later for filtering and citations.
2. Embed with KerasHub
We use a sentence-embedding preset and mean-pool the token embeddings with an attention mask:
backbone = keras_hub.models.Backbone.from_preset("bert_base_en_uncased")
preprocessor = keras_hub.models.TextClassifierPreprocessor.from_preset(
"bert_base_en_uncased", sequence_length=512
)
def embed(texts, batch_size=32):
out = []
for i in range(0, len(texts), batch_size):
batch = tf.constant(texts[i : i + batch_size])
inputs = preprocessor(batch)
seq = backbone(inputs) # (B, T, H)
mask = tf.cast(inputs["padding_mask"], seq.dtype)[..., None]
pooled = tf.reduce_sum(seq * mask, axis=1) / tf.maximum(
tf.reduce_sum(mask, axis=1), 1e-9
)
pooled = tf.math.l2_normalize(pooled, axis=-1)
out.append(pooled.numpy())
return np.concatenate(out, axis=0)
chunk_texts = [c["text"] for c in chunks]
embeddings = embed(chunk_texts).astype("float32")
Two things matter more than the choice of preset:
- Normalise. With L2-normalised vectors, dot product is cosine similarity, which is what ScaNN scores fastest.
- Use the same encoder and the same preprocessing for queries and documents. Mixing encoders, or normalising one side and not the other, produces a retriever that looks like it works on your ten test queries and collapses on real ones.
If a preset supports asymmetric prefixes ("query: " / "passage: "), use them - it is free recall. Check the KerasHub model docs for what your chosen preset expects.
3. Build a ScaNN index
import scann
searcher = (
scann.scann_ops_pybind.builder(embeddings, 10, "dot_product")
.tree(num_leaves=2000, num_leaves_to_search=100, training_sample_size=250000)
.score_ah(2, anisotropic_quantization_threshold=0.2)
.reorder(100)
.build()
)
searcher.serialize("index/v1")
Rules of thumb: num_leaves around the square root of your corpus size, num_leaves_to_search a few percent of that, and reorder set to roughly 10x your final k. Below ~50k chunks, skip the tree entirely and brute-force it with score_brute_force(); exact search on 50k vectors is a millisecond and one fewer thing to tune.
Keep a parallel list mapping row index to chunk metadata, and version the index alongside the embedding model. An index built with one encoder and queried with another returns confident nonsense - the single most common RAG bug we are called in to fix.
4. Retrieve, filter, rerank
def retrieve(question, k=8, allowed_docs=None):
q = embed([f"query: {question}"])[0]
idx, scores = searcher.search(q, final_num_neighbors=50)
hits = []
for i, s in zip(idx, scores):
meta = chunks[i]
if allowed_docs and meta["doc_id"] not in allowed_docs:
continue # tenancy / ACL filtering happens HERE, not in the prompt
hits.append({**meta, "score": float(s)})
if len(hits) == k:
break
return hits
Retrieve wide (50), filter, then cut to the k you will actually put in the prompt. Two upgrades worth their cost:
- Hybrid search. Blend the dense score with a lexical BM25 score. Dense retrieval is weak on exact identifiers - part numbers, error codes, ticket IDs - and lexical search is excellent at them.
- Cross-encoder reranking. Score the top 50 with a model that reads query and passage together, then keep the top 5. Slower per query, and usually the largest single jump in answer quality.
Never rely on the prompt for access control. If a user must not see a document, it must not enter the candidate set.
5. Generate a grounded answer
gemma = keras_hub.models.GemmaCausalLM.from_preset("gemma2_2b_en")
PROMPT = """You are a support assistant. Answer the question using ONLY the context below.
Cite sources as [1], [2] matching the numbered context blocks.
If the context does not contain the answer, reply exactly: "I don't know based on the available documents."
Context:
{context}
Question: {question}
Answer:"""
def answer(question, k=5):
hits = retrieve(question, k=k)
context = "\n\n".join(
f"[{i+1}] ({h['chunk_id']}) {h['text']}" for i, h in enumerate(hits)
)
prompt = PROMPT.format(context=context, question=question)
text = gemma.generate(prompt, max_length=1024)
return {"answer": text[len(prompt):].strip(), "sources": hits}
Details that decide whether users trust the system:
- An explicit refusal string. Without a sanctioned way to say "I don't know", the model will invent something.
- Return sources with the answer and render them as links. Users verify; that is a feature, not a failure.
- Order matters. Models attend most strongly to the beginning and end of the context. Put the best chunk first.
- Budget the context. Count tokens with the preprocessor; do not assume five chunks fit.
- A 2B-class model is enough for extract-and-cite answering when retrieval is good. Weak retrieval is not fixed by a bigger model, it is only made more expensive.
6. Evaluate - retrieval and generation separately
This is what separates a demo from a system. Build a small gold set: 50-200 real questions, each with the chunk ID(s) that contain the answer and a reference answer.
Retrieval metrics (cheap, run on every change):
def recall_at_k(gold, k=5):
hit = 0
for q, relevant in gold:
got = {h["chunk_id"] for h in retrieve(q, k=k)}
hit += bool(got & set(relevant))
return hit / len(gold)
Track Recall@5, Recall@20 and MRR. If Recall@20 is low, no amount of prompt engineering will save the answers - fix chunking, the encoder, or add hybrid search.
Generation metrics (run before each release):
- Groundedness - is every claim supported by a cited chunk?
- Answer correctness against the reference.
- Refusal accuracy - include questions your corpus genuinely cannot answer and confirm the system declines them. This one catches regressions nothing else does.
An LLM judge scoring these on a rubric is fine for regression testing, provided you calibrate it once against human labels on ~50 examples. Store every run's scores next to the index version, the encoder preset, and k, so a quality change can be traced to a config change.
7. Serving and operations
The pipeline splits cleanly in two:
- Offline (batch): chunk, embed, build and serialise the index. Rerun on a schedule or on document change. Publish as a new index version; do not mutate in place.
- Online (request path): embed the query, search, rerank, generate.
Export the query encoder as a SavedModel and put it behind TensorFlow Serving as described in Serving TensorFlow Models in Production; load the ScaNN index in the retrieval service and hot-swap it by pointing at a new directory. Cache query embeddings - support traffic is remarkably repetitive.
What to monitor:
- Retrieval score distribution. A drop in top-1 similarity means the questions have drifted from the corpus, usually before users complain. Same discipline as in Catching Data Drift Early with TFDV.
- Refusal rate. Rising = stale or incomplete corpus. Falling to near zero = the model has started guessing.
- Index freshness lag - how old is the newest indexed document?
- End-to-end p95 latency, split into embed / search / rerank / generate so you know which stage to fix.
- A thumbs-down channel with the question, retrieved chunk IDs, and generated answer logged. This is the highest-value dataset you will build, and it is what feeds your next gold set.
Where projects go wrong
Almost every failed RAG pilot we are asked to review fails for one of four reasons: chunks that cross topic boundaries, a query encoder that doesn't match the document encoder, no evaluation set so nobody can tell whether changes help, and access control implemented in the prompt instead of the retrieval filter. None of them are model problems - and all of them are cheap to fix in week one and expensive to fix in month six.
If you are scoping a grounded-answering system on your own documents, our Generative AI on the TensorFlow Stack page describes how we run these engagements, and you can get in touch for a short feasibility review of your corpus before you commit to a build.