6. Pandas for GenAI¶

Intermediate · 12 min read

Four everyday jobs: prepare RAG chunks, estimate cost, compare models, and export training data.

6.1 Documents → RAG chunks¶

import json
import pandas as pd

docs = pd.DataFrame({
    "source": ["refund-policy.md", "shipping.md"],
    "text": [
        "Refunds are allowed within 30 days of purchase. Items must be unused. "
        "Refunds go back to the original payment method within 5-7 working days.",
        "We ship across India. Delivery takes 3-5 working days. Shipping is free over ₹500.",
    ],
})

def split_sentences(text, per_chunk=2):
    """Group sentences into chunks of `per_chunk` sentences."""
    sentences = [s.strip() + "." for s in text.split(".") if s.strip()]
    return [" ".join(sentences[i:i + per_chunk]) for i in range(0, len(sentences), per_chunk)]

docs["chunk"] = docs["text"].apply(split_sentences)
chunks = docs.explode("chunk", ignore_index=True).drop(columns="text")
chunks["chunk_id"] = chunks["source"] + "#" + chunks.groupby("source").cumcount().astype(str)
chunks["n_words"] = chunks["chunk"].str.split().str.len()
print(chunks[["chunk_id", "n_words"]])
for cid, text in zip(chunks["chunk_id"], chunks["chunk"]):
    print(f"{cid}: {text}")

Output

             chunk_id  n_words
0  refund-policy.md#0       12
1  refund-policy.md#1       12
2       shipping.md#0        9
3       shipping.md#1        5
refund-policy.md#0: Refunds are allowed within 30 days of purchase. Items must be unused.
refund-policy.md#1: Refunds go back to the original payment method within 5-7 working days.
shipping.md#0: We ship across India. Delivery takes 3-5 working days.
shipping.md#1: Shipping is free over ₹500.

6.2 Clean the chunks: duplicates and tiny pieces¶

extra = pd.DataFrame({"source": ["faq.md", "faq.md"],
                      "chunk": ["Refunds are allowed within 30 days of purchase. Items must be unused.", "OK."]})
all_chunks = pd.concat([chunks, extra], ignore_index=True)
all_chunks["norm"] = all_chunks["chunk"].str.lower().str.strip()
clean = (all_chunks.drop_duplicates(subset="norm")              # same text from two files
                   .query("chunk.str.split().str.len() >= 3", engine="python")   # drop tiny chunks
                   .drop(columns="norm"))
print(len(all_chunks), "→", len(clean), "chunks")

Output

6 → 4 chunks

6.3 Estimate tokens and embedding cost¶

clean = clean.assign(est_tokens=lambda d: (d["chunk"].str.len() / 4).round().astype(int))  # ≈ 4 chars/token
usd_per_1m = 0.02                           # example embedding price — check your provider
total_tokens = clean["est_tokens"].sum()
print(clean.groupby("source")["est_tokens"].sum())
print(f"{total_tokens} tokens → ${total_tokens / 1e6 * usd_per_1m:.6f} to embed")

Output

source
refund-policy.md    35
shipping.md         21
Name: est_tokens, dtype: int64
56 tokens → $0.000001 to embed

6.4 Compare models from evaluation logs¶

logs = pd.DataFrame({
    "model":      ["gpt-4o-mini", "llama-3.1-8b", "gpt-4o"] * 3,
    "latency_ms": [820, 410, 1350, 960, 450, 1500, 790, 2900, 1280],
    "in_tokens":  [900, 900, 900, 1200, 1200, 1200, 700, 700, 700],
    "out_tokens": [210, 190, 240, 260, 230, 300, 180, 0, 220],
    "rating":     [5, 4, 5, 4, 3, 5, 5, None, 4],
    "error":      [None, None, None, None, None, None, None, "timeout", None],
})
prices = pd.DataFrame({"model": ["gpt-4o-mini", "llama-3.1-8b", "gpt-4o"],
                       "in_per_1m": [0.15, 0.05, 2.50], "out_per_1m": [0.60, 0.08, 10.00]})

logs = logs.merge(prices, on="model", how="left")
logs["cost_usd"] = (logs["in_tokens"] * logs["in_per_1m"] + logs["out_tokens"] * logs["out_per_1m"]) / 1e6

summary = logs.groupby("model").agg(
    requests=("model", "size"),
    errors=("error", "count"),                 # count() skips None → number of errors
    avg_rating=("rating", "mean"),
    p95_latency=("latency_ms", lambda s: s.quantile(0.95)),
    cost_per_1k=("cost_usd", lambda s: s.mean() * 1000),
).round(3)
print(summary.sort_values("avg_rating", ascending=False))

Output

              requests  errors  avg_rating  p95_latency  cost_per_1k
model                                                               
gpt-4o               3       0       4.667       1485.0        4.867
gpt-4o-mini          3       0       4.667        946.0        0.270
llama-3.1-8b         3       1       3.500       2655.0        0.058

6.5 Find the problems¶

slow = logs[logs["latency_ms"] > logs.groupby("model")["latency_ms"].transform("median") * 1.5]
print(slow[["model", "latency_ms", "error"]])                   # much slower than that model's normal
print(logs["error"].value_counts(dropna=True))

Output

          model  latency_ms    error
7  llama-3.1-8b        2900  timeout
error
timeout    1
Name: count, dtype: int64

6.6 Export JSONL for fine-tuning or evaluation¶

qa = pd.DataFrame({
    "question": ["How long do refunds take?", "Is shipping free?"],
    "answer":   ["Refunds reach your account within 5-7 working days.", "Yes, on orders over ₹500."],
})

def to_chat(row):
    return {"messages": [
        {"role": "system", "content": "You are a helpful support assistant."},
        {"role": "user", "content": row["question"]},
        {"role": "assistant", "content": row["answer"]},
    ]}

with open("train.jsonl", "w", encoding="utf-8") as f:
    for record in qa.apply(to_chat, axis=1):
        f.write(json.dumps(record, ensure_ascii=False) + "\n")

first = json.loads(open("train.jsonl", encoding="utf-8").readline())
print(first["messages"][1]["content"], "→", first["messages"][2]["content"])

Output

How long do refunds take? → Refunds reach your account within 5-7 working days.

Big files

For logs larger than memory, read in pieces: pd.read_csv(path, chunksize=100_000) returns an iterator of DataFrames. Save intermediate results as Parquet — it's smaller and keeps types.

Practice¶

Add a cost_inr column to logs (choose an exchange rate) and find the most expensive request.
Split docs into chunks of 1 sentence instead of 2 and compare the token estimate.

Back to: Pandas overview · Related: NumPy for GenAI