A Coding Implementation to Parsing, Analyzing, Visualizing, and Fine-Tuning Agent Reasoning Traces Using the lambda/hermes-agent-reasoning-traces Dataset

!pip -q install -U datasets pandas matplotlib seaborn transformers accelerate trl

import json, re, random, textwrap
from collections import Counter, defaultdict
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datasets import load_dataset, concatenate_datasets

random.seed(0)

CONFIG = “kimi”
ds = load_dataset(“lambda/hermes-agent-reasoning-traces”, CONFIG, split=”train”)
print(ds)
print(“Config:”, CONFIG, “| Fields:”, ds.column_names)
print(“Categories:”, sorted(set(ds[“category”])))

COMPARE_BOTH = False
if COMPARE_BOTH:
ds_kimi = load_dataset(“lambda/hermes-agent-reasoning-traces”, “kimi”, split=”train”)
ds_glm = load_dataset(“lambda/hermes-agent-reasoning-traces”, “glm-5.1″, split=”train”)
ds_kimi = ds_kimi.add_column(“source”, [“kimi”] * len(ds_kimi))
ds_glm = ds_glm.add_column(“source”, [“glm-5.1”] * len(ds_glm))
ds = concatenate_datasets([ds_kimi, ds_glm]).shuffle(seed=0)
print(“Combined:”, ds, “→ counts:”, Counter(ds[“source”]))

sample = ds[0]
print(“\n=== Sample 0 ===”)
print(“id :”, sample[“id”])
print(“category :”, sample[“category”], “/”, sample[“subcategory”])
print(“task :”, sample[“task”])
print(“turns :”, len(sample[“conversations”]))
print(“system[0] :”, sample[“conversations”][0][“value”][:220], “…\n”)

THINK_RE = re.compile(r”<think>(.*?)</think>”, re.DOTALL)
TOOL_CALL_RE = re.compile(r”<tool_call>\s*(\{.*?\})\s*</tool_call>”, re.DOTALL)
TOOL_RESP_RE = re.compile(r”<tool_response>\s*(.*?)\s*</tool_response>”, re.DOTALL)

def parse_assistant(value: str) -> dict:
thoughts = [t.strip() for t in THINK_RE.findall(value)]
calls = []
for raw in TOOL_CALL_RE.findall(value):
try:
calls.append(json.loads(raw))
except json.JSONDecodeError:
calls.append({“name”: “<malformed>”, “arguments”: {}})
final = TOOL_CALL_RE.sub(“”, THINK_RE.sub(“”, value)).strip()
return {“thoughts”: thoughts, “tool_calls”: calls, “final”: final}

def parse_tool(value: str):
raw = TOOL_RESP_RE.search(value)
if not raw: return {“raw”: value}
body = raw.group(1)
try: return json.loads(body)
except: return {“raw”: body}

first_gpt = next(t for t in sample[“conversations”] if t[“from”] == “gpt”)
p = parse_assistant(first_gpt[“value”])
print(“Thought preview :”, (p[“thoughts”][0][:160] + “…”) if p[“thoughts”] else “(none)”)
print(“Tool calls :”, [(c.get(“name”), list(c.get(“arguments”, {}).keys())) for c in p[“tool_calls”]])

N = 3000
sub = ds.select(range(min(N, len(ds))))

tool_calls = Counter()
parallel_widths = Counter()
thoughts_per_turn = []
calls_per_traj = []
errors_per_traj = []
turns_per_traj = []
cat_counts = Counter()

for ex in sub:
cat_counts[ex[“category”]] += 1
n_calls = n_err = 0
turns_per_traj.append(len(ex[“conversations”]))
for t in ex[“conversations”]:
if t[“from”] == “gpt”:
p = parse_assistant(t[“value”])
thoughts_per_turn.append(len(p[“thoughts”]))
if p[“tool_calls”]:
parallel_widths[len(p[“tool_calls”])] += 1
for c in p[“tool_calls”]:
tool_calls[c.get(“name”, “<unknown>”)] += 1
n_calls += len(p[“tool_calls”])
elif t[“from”] == “tool”:
r = parse_tool(t[“value”])
blob = json.dumps(r).lower()
if “error” in blob or ‘”exit_code”: 1’ in blob or “traceback” in blob:
n_err += 1
calls_per_traj.append(n_calls)
errors_per_traj.append(n_err)

print(f”\nScanned {len(sub)} trajectories”)
print(f”Avg turns/traj : {np.mean(turns_per_traj):.1f}”)
print(f”Avg tool calls/traj : {np.mean(calls_per_traj):.1f}”)
print(f”% with >=1 error : {100*np.mean([e>0 for e in errors_per_traj]):.1f}%”)
print(f”% parallel turns : {100*sum(v for k,v in parallel_widths.items() if k>1)/max(1,sum(parallel_widths.values())):.1f}%”)
print(“Top 10 tools :”, tool_calls.most_common(10))

fig, axes = plt.subplots(2, 2, figsize=(13, 9))

top = tool_calls.most_common(15)
axes[0,0].barh([t for t,_ in top][::-1], [c for _,c in top][::-1], color=”teal”)
axes[0,0].set_title(“Top 15 tools by call volume”)
axes[0,0].set_xlabel(“calls”)

ks = sorted(parallel_widths)
axes[0,1].bar([str(k) for k in ks], [parallel_widths[k] for k in ks], color=”coral”)
axes[0,1].set_title(“Tool-calls per assistant turn (parallel width)”)
axes[0,1].set_xlabel(“# tool calls in one turn”); axes[0,1].set_ylabel(“count”)
axes[0,1].set_yscale(“log”)

axes[1,0].hist(turns_per_traj, bins=40, color=”steelblue”)
axes[1,0].set_title(“Conversation length”); axes[1,0].set_xlabel(“turns”)

cats, vals = zip(*cat_counts.most_common())
axes[1,1].pie(vals, labels=cats, autopct=”%1.0f%%”, startangle=90)
axes[1,1].set_title(“Category distribution”)

plt.tight_layout(); plt.show()

def render_trace(ex, max_chars=350):
print(f”\n{‘=’*72}\nTASK [{ex[‘category’]} / {ex[‘subcategory’]}]: {ex[‘task’]}\n{‘=’*72}”)
for t in ex[“conversations”]:
role = t[“from”]
if role == “system”:
continue
if role == “human”:
print(f”\n[USER]\n{textwrap.shorten(t[‘value’], 600)}”)
elif role == “gpt”:
p = parse_assistant(t[“value”])
for th in p[“thoughts”]:
print(f”\n[THINK]\n{textwrap.shorten(th, max_chars)}”)
for c in p[“tool_calls”]:
args = json.dumps(c.get(“arguments”, {}))[:200]
print(f”[CALL] {c.get(‘name’)}({args})”)
if p[“final”]:
print(f”\n[ANSWER]\n{textwrap.shorten(p[‘final’], max_chars)}”)
elif role == “tool”:
print(f”[TOOL_RESPONSE] {textwrap.shorten(t[‘value’], 220)}”)
print(“=”*72)

idx = int(np.argmin(np.abs(np.array(turns_per_traj) – 10)))
render_trace(sub[idx])

def get_tool_schemas(ex):
try: return json.loads(ex[“tools”])
except: return []

schemas = get_tool_schemas(sample)
print(f”\nSample 0 has {len(schemas)} tools available”)
for s in schemas[:3]:
fn = s.get(“function”, {})
print(” -“, fn.get(“name”), “—”, (fn.get(“description”) or “”)[:80])

ROLE_MAP = {“system”: “system”, “human”: “user”, “gpt”: “assistant”, “tool”: “tool”}

def to_openai_messages(conv):
return [{“role”: ROLE_MAP[t[“from”]], “content”: t[“value”]} for t in conv]

example_msgs = to_openai_messages(sample[“conversations”])
print(“\nFirst 2 OpenAI messages:”)
for m in example_msgs[:2]:
print(” “, m[“role”], “→”, m[“content”][:120].replace(“\n”, ” “), “…”)

from transformers import AutoTokenizer
TOK_ID = “Qwen/Qwen2.5-0.5B-Instruct”
tok = AutoTokenizer.from_pretrained(TOK_ID)

def build_masked(conv, tokenizer, max_len=2048):
msgs = to_openai_messages(conv)
for m in msgs:
if m[“role”] == “tool”:
m[“role”] = “user”
m[“content”] = “[TOOL OUTPUT]\n” + m[“content”]
input_ids, labels = [], []
for m in msgs:
text = tokenizer.apply_chat_template([m], tokenize=False, add_generation_prompt=False)
ids = tokenizer.encode(text, add_special_tokens=False)
input_ids.extend(ids)
labels.extend(ids if m[“role”] == “assistant” else [-100] * len(ids))
return input_ids[:max_len], labels[:max_len]

ids, lbls = build_masked(sample[“conversations”], tok)
trainable = sum(1 for x in lbls if x != -100)
print(f”\nTokenized example: {len(ids)} tokens, {trainable} trainable ({100*trainable/len(ids):.1f}%)”)

think_lens, call_lens, ans_lens = [], [], []
for ex in sub.select(range(min(500, len(sub)))):
for t in ex[“conversations”]:
if t[“from”] != “gpt”: continue
p = parse_assistant(t[“value”])
for th in p[“thoughts”]: think_lens.append(len(th))
for c in p[“tool_calls”]: call_lens.append(len(json.dumps(c)))
if p[“final”]: ans_lens.append(len(p[“final”]))

plt.figure(figsize=(10,4))
plt.hist([think_lens, call_lens, ans_lens], bins=40, log=True,
label=[“<think>”, “<tool_call>”, “final answer”], stacked=False)
plt.legend(); plt.xlabel(“characters”); plt.title(“Length distributions (log y)”)
plt.tight_layout(); plt.show()

class TraceReplayer:
def __init__(self, ex):
self.ex = ex
self.steps = []
pending = None
for t in ex[“conversations”]:
if t[“from”] == “gpt”:
if pending: self.steps.append(pending)
pending = {“think”: parse_assistant(t[“value”]), “responses”: []}
elif t[“from”] == “tool” and pending:
pending[“responses”].append(parse_tool(t[“value”]))
if pending: self.steps.append(pending)
def __len__(self): return len(self.steps)
def play(self, i):
s = self.steps[i]
print(f”\n── Step {i+1}/{len(self)} ──”)
for th in s[“think”][“thoughts”]:
print(f”💭 {textwrap.shorten(th, 280)}”)
for c in s[“think”][“tool_calls”]:
print(f”⚙️ {c.get(‘name’)}({json.dumps(c.get(‘arguments’, {}))[:140]})”)
for r in s[“responses”]:
print(f”📥 {textwrap.shorten(json.dumps(r), 200)}”)
if s[“think”][“final”]:
print(f”💬 {textwrap.shorten(s[‘think’][‘final’], 200)}”)

rp = TraceReplayer(sample)
for i in range(min(3, len(rp))):
rp.play(i)

TRAIN = False
if TRAIN:
import torch
from transformers import AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig

train_subset = ds.select(range(200))

def to_text(batch):
msgs = to_openai_messages(batch[“conversations”])
for m in msgs:
if m[“role”] == “tool”:
m[“role”] = “user”; m[“content”] = “[TOOL]\n” + m[“content”]
batch[“text”] = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False)
return batch

train_subset = train_subset.map(to_text)

model = AutoModelForCausalLM.from_pretrained(
TOK_ID,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
device_map=”auto” if torch.cuda.is_available() else None,
)

cfg = SFTConfig(
output_dir=”hermes-sft-demo”,
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
max_steps=20,
learning_rate=2e-5,
logging_steps=2,
max_seq_length=1024,
dataset_text_field=”text”,
report_to=”none”,
fp16=torch.cuda.is_available(),
)
SFTTrainer(model=model, args=cfg, train_dataset=train_subset, processing_class=tok).train()
print(“Fine-tune demo finished.”)

print(“\n✅ Tutorial complete. You now have parsers, analytics, plots, a replayer, ”
“tokenized + label-masked SFT examples, and an optional training hook.”)