AI Data Pipeline Patterns

Comprehensions, generators and itertools applied to the code AI engineers write every day: prompt construction, chat messages, embedding batches, retrieval post-processing, tool calling, token budgeting, dataset streaming, evaluation and training loops. Each pattern names the shape (map, filter, flatten, group, reduce) and when to keep a loop instead.

Prompt and message construction
Chunking and text preparation
Embedding batches
Vector store records
Retrieval post-processing
Tool calling
Token counting and context budgets
Streaming responses
Lazy dataset streaming
Evaluation
Training and inference loops
Async fan-out
Pattern index

Prompt and message construction

python

SYSTEM = "You are a concise assistant. Cite sources as [n]."

def build_messages(question: str, history: list[tuple[str, str]], hits: list) -> list[dict]:
    context = "\n\n".join(f"[{i}] {h.text}" for i, h in enumerate(hits, start=1))
    return [
        {"role": "system", "content": SYSTEM},
        *({"role": role, "content": text} for role, text in history),
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"},
    ]

Few-shot examples from a dataset, skipping ones that are too long:

python

shots = [
    f"Q: {ex['question']}\nA: {ex['answer']}"
    for ex in examples
    if len(ex["question"]) + len(ex["answer"]) < 600
][:5]
few_shot_block = "\n\n".join(shots)

Prompt list for batch inference:

python

prompts = [TEMPLATE.format(question=q.strip(), lang=lang) for q in questions if q.strip()]

Chat-format conversion between providers (OpenAI-style to a list of role/content pairs):

python

pairs = [(m["role"], m["content"]) for m in messages if m["role"] in {"user", "assistant"}]

Stripping tool messages or None contents from history before resending:

python

history = [m for m in history if m["role"] != "tool" and m.get("content")]

Chunking and text preparation

python

paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]

# fixed-size token windows with overlap
def chunk_tokens(tokens: list[int], size: int = 512, overlap: int = 64) -> list[list[int]]:
    step = size - overlap
    return [tokens[i:i + size] for i in range(0, max(len(tokens) - overlap, 1), step)]

# flatten chunks from many documents, keeping provenance
chunks = [
    {"doc_id": doc.id, "chunk_idx": i, "text": piece}
    for doc in docs
    for i, piece in enumerate(splitter.split_text(doc.text))
]

Deduplicate near-identical chunks (keeps first-seen order; a loop because it carries state):

python

import hashlib

def dedupe(chunks):
    seen = set()
    for c in chunks:
        h = hashlib.sha1(" ".join(c["text"].split()).lower().encode()).hexdigest()
        if h not in seen:
            seen.add(h)
            yield c

Embedding batches

Never embed one text per request, and never send all texts in one request. Batch lazily.

python

from itertools import batched   # 3.12+; backport in generators-and-itertools.md

def embed_all(texts, client, model="text-embedding-3-small", batch_size=128):
    vectors = []
    for batch in batched(texts, batch_size):
        resp = client.embeddings.create(model=model, input=list(batch))
        vectors.extend(d.embedding for d in sorted(resp.data, key=lambda d: d.index))
    return vectors

This is a loop, not a comprehension, because each iteration does I/O and may need retries. The inner extend with a generator is idiomatic.

With rate-limit retry (keep it in a loop; retries are control flow):

python

import random
import time

def embed_with_retry(batch, client, model, attempts=5):
    for attempt in range(attempts):
        try:
            return [d.embedding for d in client.embeddings.create(model=model, input=list(batch)).data]
        except RateLimitError:
            time.sleep(min(2 ** attempt, 30) + random.random())
    raise RuntimeError("embedding failed after retries")

vectors = [v for batch in batched(texts, 128) for v in embed_with_retry(batch, client, MODEL)]

Token-budgeted batches (APIs cap tokens per request) need a generator function: see batch_by_tokens in generators-and-itertools.md.

Vector store records

python

ids = [c["id"] for c in chunks]
texts = [c["text"] for c in chunks]
metadatas = [{k: v for k, v in c.items() if k not in {"text", "embedding"}} for c in chunks]

# id -> embedding lookup
vec_by_id = {c["id"]: c["embedding"] for c in chunks}

# upsert in batches; zip keeps parallel lists aligned (strict=True: 3.10+)
for batch in batched(zip(ids, vectors, metadatas, strict=True), 100):
    index.upsert(vectors=[{"id": i, "values": v, "metadata": m} for i, v, m in batch])

Convert to a numpy matrix once for local search:

python

import numpy as np

order = list(vec_by_id)                                     # stable id order
matrix = np.array([vec_by_id[i] for i in order], dtype=np.float32)

Retrieval post-processing

python

# keep hits above a threshold, best first
relevant = sorted((h for h in hits if h.score >= 0.75), key=lambda h: h.score, reverse=True)

# first good hit or None
best = next((h for h in hits if h.score >= 0.9), None)

# unique sources, in rank order
sources = list(dict.fromkeys(h.metadata["source"] for h in relevant))

# group hits per document and keep the best chunk of each
from collections import defaultdict

by_doc = defaultdict(list)
for h in relevant:
    by_doc[h.metadata["doc_id"]].append(h)
best_per_doc = [max(group, key=lambda h: h.score) for group in by_doc.values()]

Reciprocal rank fusion of several retrievers (dict accumulation with a loop, then a sorted comprehension):

python

def rrf(result_lists, k=60):
    scores = defaultdict(float)
    for results in result_lists:
        for rank, doc_id in enumerate(results, start=1):
            scores[doc_id] += 1.0 / (k + rank)
    return [doc_id for doc_id, _ in sorted(scores.items(), key=lambda kv: kv[1], reverse=True)]

Score normalisation:

python

lo, hi = min(scores), max(scores)
norm = [(s - lo) / (hi - lo) if hi > lo else 0.0 for s in scores]

Tool calling

Tool schemas from Python functions:

python

import inspect

def tool_schema(fn):
    params = inspect.signature(fn).parameters
    return {
        "type": "function",
        "function": {
            "name": fn.__name__,
            "description": inspect.getdoc(fn) or "",
            "parameters": {
                "type": "object",
                "properties": {name: {"type": "string"} for name in params},
                "required": [n for n, p in params.items() if p.default is inspect.Parameter.empty],
            },
        },
    }

TOOLS = [search_docs, get_weather, run_sql]
schemas = [tool_schema(fn) for fn in TOOLS]
registry = {fn.__name__: fn for fn in TOOLS}

Flatten tool calls across choices or turns, tolerating None:

python

calls = [call for choice in resp.choices for call in (choice.message.tool_calls or [])]
names = {c.function.name for c in calls}

Execute tool calls: a loop, because each one is an action with error handling:

python

import json

tool_messages = []
for call in calls:
    fn = registry.get(call.function.name)
    try:
        result = fn(**json.loads(call.function.arguments)) if fn else {"error": "unknown tool"}
    except Exception as exc:   # report the failure back to the model
        result = {"error": str(exc)}
    tool_messages.append({"role": "tool", "tool_call_id": call.id, "content": json.dumps(result)})

Anthropic-style content blocks:

python

text = "".join(b.text for b in message.content if b.type == "text")
tool_uses = [b for b in message.content if b.type == "tool_use"]

Token counting and context budgets

python

from collections import Counter

total = sum(len(enc.encode(m["content"])) for m in messages)
per_role = Counter()
for m in messages:
    per_role[m["role"]] += len(enc.encode(m["content"]))
vocab_freq = Counter(tok for doc in corpus for tok in enc.encode(doc))

Count once, reuse (walrus):

python

sized = [(m, n) for m in messages if (n := len(enc.encode(m["content"]))) > 0]

Trim history to a budget, keeping the most recent messages (needs early exit, so a loop):

python

def fit_history(messages, budget, count):
    kept, used = [], 0
    for m in reversed(messages):
        n = count(m["content"])
        if used + n > budget:
            break
        kept.append(m)
        used += n
    return kept[::-1]

Cumulative token offsets for chunk citation spans:

python

from itertools import accumulate, pairwise

lengths = [len(enc.encode(c)) for c in chunks]
spans = list(pairwise(accumulate(lengths, initial=0)))   # [(0, 120), (120, 460), ...]  3.10+

Streaming responses

python

def stream_text(stream):
    for event in stream:
        delta = event.choices[0].delta.content
        if delta:
            yield delta

answer = "".join(stream_text(stream))

Async streaming with an async comprehension (inside async def):

python

parts = [e.choices[0].delta.content async for e in stream if e.choices and e.choices[0].delta.content]
answer = "".join(parts)

Group streamed events by message id when a stream interleaves several (already ordered by id runs):

python

from itertools import groupby
messages = {mid: "".join(e.text for e in evs) for mid, evs in groupby(events, key=lambda e: e.message_id)}

Lazy dataset streaming

Keep memory flat regardless of dataset size by chaining generators:

python

import json
from itertools import islice

def read_jsonl(path):
    with open(path, encoding="utf-8") as fh:
        for line in fh:
            if line.strip():
                yield json.loads(line)

records = read_jsonl("train.jsonl")
valid = (r for r in records if r.get("text") and r.get("label") is not None)
formatted = ({"prompt": PROMPT.format(text=r["text"]), "completion": r["label"]} for r in valid)
preview = list(islice(formatted, 3))   # inspect without reading the whole file

Many shards as one stream:

python

from itertools import chain
from pathlib import Path

shards = sorted(Path("data/").glob("part-*.jsonl"))
all_records = chain.from_iterable(read_jsonl(p) for p in shards)

Deterministic train/validation split without loading everything (hash-based):

python

import zlib

def is_val(record_id: str, pct: int = 5) -> bool:
    return zlib.crc32(record_id.encode()) % 100 < pct

train = (r for r in read_jsonl(path) if not is_val(r["id"]))
val = (r for r in read_jsonl(path) if is_val(r["id"]))

Hugging Face datasets: prefer ds.map(fn, batched=True) and ds.filter(...) (Arrow-backed, multiprocess) over Python comprehensions over rows. Use comprehensions only on the small batches inside fn:

python

def tokenize(batch):
    return tokenizer([t.strip() for t in batch["text"]], truncation=True, max_length=512)

ds = ds.map(tokenize, batched=True)

Evaluation

python

preds = [normalize(p) for p in predictions]
golds = [normalize(g) for g in references]
exact = sum(p == g for p, g in zip(preds, golds, strict=True)) / len(golds)

# per-category accuracy
from collections import defaultdict
buckets = defaultdict(list)
for case, p in zip(cases, preds):
    buckets[case["category"]].append(p == normalize(case["answer"]))
accuracy = {cat: sum(v) / len(v) for cat, v in buckets.items()}

failures = [c["id"] for c, p in zip(cases, preds) if p != normalize(c["answer"])]
all_passed = all(p == g for p, g in zip(preds, golds))

Hyper-parameter or prompt-variant grid:

python

from itertools import product

runs = [
    {"model": m, "temperature": t, "prompt": name}
    for m, t, name in product(["small", "large"], [0.0, 0.7], PROMPTS)
]

Training and inference loops

Training loops are loops: they have side effects (optimizer steps, logging, checkpointing). Comprehensions appear around them.

python

import torch

param_groups = [
    {"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in ("bias", "norm"))],
     "weight_decay": 0.01},
    {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in ("bias", "norm"))],
     "weight_decay": 0.0},
]
optimizer = torch.optim.AdamW(param_groups, lr=3e-5)

n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
frozen = [n for n, p in model.named_parameters() if not p.requires_grad]

batch = {k: v.to(device) for k, v in batch.items()}           # move a batch dict to GPU
metrics = {k: v.item() for k, v in loss_parts.items()}          # one .item() per scalar, not per element

Collate function with padding:

python

def collate(samples, pad_id=0):
    max_len = max(len(s["input_ids"]) for s in samples)
    input_ids = [s["input_ids"] + [pad_id] * (max_len - len(s["input_ids"])) for s in samples]
    attention = [[1] * len(s["input_ids"]) + [0] * (max_len - len(s["input_ids"])) for s in samples]
    return {"input_ids": torch.tensor(input_ids), "attention_mask": torch.tensor(attention)}

Batched inference over a generator of inputs:

python

@torch.inference_mode()
def predict(texts, batch_size=32):
    for batch in batched(texts, batch_size):
        enc = tokenizer(list(batch), padding=True, truncation=True, return_tensors="pt").to(device)
        yield from model(**enc).logits.argmax(dim=-1).tolist()

labels = [id2label[i] for i in predict(texts)]

Async fan-out

A comprehension with await runs sequentially. Build coroutines with a generator and gather them, with a semaphore for rate limits:

python

import asyncio

async def summarize_all(docs, client, limit=10):
    sem = asyncio.Semaphore(limit)

    async def one(doc):
        async with sem:
            resp = await client.messages.create(model=MODEL, max_tokens=300,
                                                messages=[{"role": "user", "content": f"Summarize:\n{doc}"}])
            return resp.content[0].text

    return await asyncio.gather(*(one(d) for d in docs))

Python 3.11+ TaskGroup (cancels siblings on failure):

python

async def fetch_all(urls, fetch):
    async with asyncio.TaskGroup() as tg:
        tasks = [tg.create_task(fetch(u)) for u in urls]
    return [t.result() for t in tasks]

Keep partial results when some calls fail:

python

results = await asyncio.gather(*(one(d) for d in docs), return_exceptions=True)
ok = [r for r in results if not isinstance(r, BaseException)]
errors = [r for r in results if isinstance(r, BaseException)]

Pattern index

Task	Shape	Tool
prompts from questions	map + filter	list comprehension
chat messages with history	map + unpack	list with `*(... for ...)`
chunk documents	flatten with provenance	two-`for` comprehension with `enumerate`
dedupe chunks	stateful filter	generator function with a `seen` set
embedding requests	batch + I/O	`batched` in a loop
id -> vector	key/value map	dict comprehension
filter hits by score	filter	comprehension / generator into `sorted`
first good hit	search	`next(gen, None)`
unique sources in order	dedupe	`dict.fromkeys`
group hits by doc	group	`defaultdict(list)` loop
flatten tool calls	flatten	two-`for` comprehension or `chain.from_iterable`
execute tool calls	side effects	for loop
token totals	reduce	`sum(generator)`
token frequencies	count	`Counter(generator)`
context budget	early exit	for loop with `break`
JSONL dataset	stream	generator functions + `islice`
eval accuracy	reduce over zip	`sum(p == g for ...)`
optimizer param groups	filter	list comprehension
concurrent LLM calls	fan-out	`asyncio.gather(*(coro for ...))`

Contents

Prompt and message construction

Chunking and text preparation

Embedding batches

Vector store records

Retrieval post-processing

Tool calling

Token counting and context budgets

Streaming responses

Lazy dataset streaming

Evaluation

Training and inference loops

Async fan-out

Pattern index

python-beginner-concept-tutor

python-context-manager-writer

python-dataclass-converter

python-exception-handling-designer

python-idiomatic-code-reviewer

python-type-hints-annotator

AI Data Pipeline Patterns

#Contents

#Prompt and message construction

#Chunking and text preparation

#Embedding batches

#Vector store records

#Retrieval post-processing

#Tool calling

#Token counting and context budgets

#Streaming responses

#Lazy dataset streaming

#Evaluation

#Training and inference loops

#Async fan-out

#Pattern index

python-beginner-concept-tutor

python-context-manager-writer

python-dataclass-converter

python-exception-handling-designer

python-idiomatic-code-reviewer

python-type-hints-annotator

Contents

Prompt and message construction

Chunking and text preparation

Embedding batches

Vector store records

Retrieval post-processing

Tool calling

Token counting and context budgets

Streaming responses

Lazy dataset streaming

Evaluation

Training and inference loops

Async fan-out

Pattern index