Spaces:

dipsha
/

recruiter

Sleeping

Upload folder using huggingface_hub

a8db748 verified 6 months ago

9.54 kB

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
plot_results.py — generate training result charts for the submission README.
Produces a 3-subplot figure saved to plots/results.png:
  Top:    Episode return curve (x=episode, y=reward), dashed baseline, curriculum tier markers
  Middle: nDCG@5 before vs after training, per difficulty tier
  Bottom: Fakes-in-top-5 before vs after training, per difficulty tier
Usage:
    python -m recruiter.training.plot_results \\
        --baseline baseline_random.json \\
        --trained   trained_eval.json \\
        --training  training_log.jsonl \\
        --output    plots/results.png
    # Minimal (just training curve, no before/after eval):
    python -m recruiter.training.plot_results --training training_log.jsonl
Input formats:
    baseline/trained JSON: output of baseline_eval.py
    training_log JSONL: one JSON object per update, e.g.:
        {"update": 0, "episode": 4, "mean_reward": -3.2, "difficulty": 1}
All files are optional — subplots are skipped if data is missing.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Optional


def _load_jsonl(path: str) -> list[dict]:
    records = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def _load_json(path: str) -> dict:
    with open(path) as f:
        return json.load(f)


def _rolling_mean(values: list[float], window: int = 20) -> list[float]:
    out = []
    for i, v in enumerate(values):
        start = max(0, i - window + 1)
        out.append(sum(values[start : i + 1]) / (i - start + 1))
    return out


def _tier_stats(eval_data: dict) -> dict[int, dict]:
    """Extract per-tier mean_ndcg and mean_fakes_in_top5 from baseline_eval output."""
    per_ep = eval_data.get("per_episode", [])
    tier_buckets: dict[int, list[dict]] = {}
    for ep in per_ep:
        d = ep.get("difficulty", 1)
        tier_buckets.setdefault(d, []).append(ep)

    result = {}
    for tier, eps in tier_buckets.items():
        ndcgs = [e.get("ndcg", 0.0) for e in eps]
        fakes = [e.get("fakes_in_top5", 0) for e in eps]
        result[tier] = {
            "mean_ndcg": sum(ndcgs) / len(ndcgs),
            "mean_fakes": sum(fakes) / len(fakes),
            "n": len(eps),
        }
    return result


def plot(
    training_log: Optional[str] = None,
    baseline_path: Optional[str] = None,
    trained_path: Optional[str] = None,
    output: str = "plots/results.png",
    wandb_url: Optional[str] = None,
) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import matplotlib.ticker as ticker
    import numpy as np

    Path(output).parent.mkdir(parents=True, exist_ok=True)

    # ── Determine number of subplots ──────────────────────────────────────────
    has_training = training_log is not None and Path(training_log).exists()
    has_before_after = (
        baseline_path is not None and Path(baseline_path).exists()
        and trained_path is not None and Path(trained_path).exists()
    )

    n_plots = 1 + (2 if has_before_after else 0)
    fig, axes = plt.subplots(n_plots, 1, figsize=(10, 4 * n_plots))
    if n_plots == 1:
        axes = [axes]

    fig.suptitle("Recruiter Agent Training Results", fontsize=14, fontweight="bold", y=0.98)

    # ── Subplot 1: Training curve ─────────────────────────────────────────────
    ax = axes[0]
    ax.set_title("Episode Return During Training")
    ax.set_xlabel("Episode")
    ax.set_ylabel("Return (reward)")

    if has_training:
        records = _load_jsonl(training_log)
        episodes = [r.get("episode", i * 4) for i, r in enumerate(records)]
        # Support both "mean_return" (train_grpo.py key) and "mean_reward" (legacy)
        rewards = [r.get("mean_terminal_return", r.get("mean_return", r.get("mean_reward", 0.0))) for r in records]
        smooth = _rolling_mean(rewards, window=10)

        ax.plot(episodes, rewards, alpha=0.3, color="steelblue", linewidth=0.8, label="raw")
        ax.plot(episodes, smooth, color="steelblue", linewidth=2.0, label="smoothed (win=10)")

        # Curriculum tier markers
        difficulties = [r.get("difficulty", 1) for r in records]
        prev_d = difficulties[0] if difficulties else 1
        for i, d in enumerate(difficulties):
            if d != prev_d:
                ax.axvline(x=episodes[i], color="orange", linestyle="--", alpha=0.6, linewidth=1.2)
                ax.text(episodes[i], ax.get_ylim()[0] if ax.get_ylim()[0] != 0 else min(rewards),
                        f"T{d}", fontsize=8, color="darkorange", ha="left")
                prev_d = d
    else:
        ax.text(0.5, 0.5, "No training log provided", transform=ax.transAxes,
                ha="center", va="center", fontsize=12, color="grey")

    # Baseline reference line
    if baseline_path and Path(baseline_path).exists():
        baseline = _load_json(baseline_path)
        mean_base = baseline.get("summary", {}).get("mean_reward", None)
        if mean_base is not None:
            ax.axhline(y=mean_base, color="red", linestyle="--", linewidth=1.5,
                       label=f"random baseline ({mean_base:.2f})")

    ax.legend(loc="lower right", fontsize=9)
    ax.grid(True, alpha=0.3)

    if has_before_after:
        baseline_data = _load_json(baseline_path)
        trained_data = _load_json(trained_path)
        before_stats = _tier_stats(baseline_data)
        after_stats = _tier_stats(trained_data)

        all_tiers = sorted(set(before_stats) | set(after_stats))
        tier_labels = [f"T{t}" for t in all_tiers]
        x = np.arange(len(all_tiers))
        width = 0.35

        # ── Subplot 2: nDCG@5 per tier ────────────────────────────────────────
        ax2 = axes[1]
        ax2.set_title("nDCG@5 Before vs After Training (per difficulty tier)")
        ax2.set_ylabel("nDCG@5")
        ax2.set_ylim(0, 1.05)

        before_ndcg = [before_stats.get(t, {}).get("mean_ndcg", 0.0) for t in all_tiers]
        after_ndcg = [after_stats.get(t, {}).get("mean_ndcg", 0.0) for t in all_tiers]

        ax2.bar(x - width / 2, before_ndcg, width, label="Before (zero-shot)", color="salmon", alpha=0.8)
        ax2.bar(x + width / 2, after_ndcg, width, label="After (fine-tuned)", color="steelblue", alpha=0.8)
        ax2.set_xticks(x)
        ax2.set_xticklabels(tier_labels)
        ax2.legend(fontsize=9)
        ax2.grid(True, alpha=0.3, axis="y")

        # Annotate delta
        for i, (b, a) in enumerate(zip(before_ndcg, after_ndcg)):
            delta = a - b
            color = "green" if delta >= 0 else "red"
            ax2.text(x[i] + width / 2, a + 0.02, f"{delta:+.2f}", ha="center",
                     fontsize=8, color=color, fontweight="bold")

        # ── Subplot 3: Fakes-in-top-5 per tier ───────────────────────────────
        ax3 = axes[2]
        ax3.set_title("Fakes in Top-5 Before vs After Training (per difficulty tier)")
        ax3.set_ylabel("Mean fakes in top-5 (lower is better)")

        before_fakes = [before_stats.get(t, {}).get("mean_fakes", 0.0) for t in all_tiers]
        after_fakes = [after_stats.get(t, {}).get("mean_fakes", 0.0) for t in all_tiers]

        ax3.bar(x - width / 2, before_fakes, width, label="Before (zero-shot)", color="salmon", alpha=0.8)
        ax3.bar(x + width / 2, after_fakes, width, label="After (fine-tuned)", color="steelblue", alpha=0.8)
        ax3.set_xticks(x)
        ax3.set_xticklabels(tier_labels)
        ax3.legend(fontsize=9)
        ax3.grid(True, alpha=0.3, axis="y")

        # Annotate delta
        for i, (b, a) in enumerate(zip(before_fakes, after_fakes)):
            delta = a - b
            color = "green" if delta <= 0 else "red"
            ax3.text(x[i] + width / 2, a + 0.02, f"{delta:+.2f}", ha="center",
                     fontsize=8, color=color, fontweight="bold")

    if wandb_url:
        fig.text(0.5, 0.01, f"W&B run: {wandb_url}", ha="center", fontsize=8, color="grey")

    plt.tight_layout(rect=[0, 0.02, 1, 0.97])
    plt.savefig(output, dpi=150, bbox_inches="tight")
    print(f"Saved plot to {output}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Plot recruiter training results")
    parser.add_argument("--training", default=None, help="Path to training_log.jsonl")
    parser.add_argument("--baseline", default=None, help="Path to baseline_eval output JSON")
    parser.add_argument("--trained", default=None, help="Path to post-training eval JSON")
    parser.add_argument("--output", default="plots/results.png", help="Output PNG path")
    parser.add_argument("--wandb-url", default=None, help="W&B run URL to embed in figure")
    args = parser.parse_args()

    plot(
        training_log=args.training,
        baseline_path=args.baseline,
        trained_path=args.trained,
        output=args.output,
        wandb_url=args.wandb_url,
    )


if __name__ == "__main__":
    main()

	# Copyright (c) Meta Platforms, Inc. and affiliates.
	# All rights reserved.
	#
	# This source code is licensed under the BSD-style license found in the
	# LICENSE file in the root directory of this source tree.

	"""
	plot_results.py — generate training result charts for the submission README.

	Produces a 3-subplot figure saved to plots/results.png:
	Top: Episode return curve (x=episode, y=reward), dashed baseline, curriculum tier markers
	Middle: nDCG@5 before vs after training, per difficulty tier
	Bottom: Fakes-in-top-5 before vs after training, per difficulty tier

	Usage:
	python -m recruiter.training.plot_results \\
	--baseline baseline_random.json \\
	--trained trained_eval.json \\
	--training training_log.jsonl \\
	--output plots/results.png

	# Minimal (just training curve, no before/after eval):
	python -m recruiter.training.plot_results --training training_log.jsonl

	Input formats:
	baseline/trained JSON: output of baseline_eval.py
	training_log JSONL: one JSON object per update, e.g.:
	{"update": 0, "episode": 4, "mean_reward": -3.2, "difficulty": 1}

	All files are optional — subplots are skipped if data is missing.
	"""

	from __future__ import annotations

	import argparse
	import json
	import os
	from pathlib import Path
	from typing import Optional


	def _load_jsonl(path: str) -> list[dict]:
	records = []
	with open(path) as f:
	for line in f:
	line = line.strip()
	if line:
	records.append(json.loads(line))
	return records


	def _load_json(path: str) -> dict:
	with open(path) as f:
	return json.load(f)


	def _rolling_mean(values: list[float], window: int = 20) -> list[float]:
	out = []
	for i, v in enumerate(values):
	start = max(0, i - window + 1)
	out.append(sum(values[start : i + 1]) / (i - start + 1))
	return out


	def _tier_stats(eval_data: dict) -> dict[int, dict]:
	"""Extract per-tier mean_ndcg and mean_fakes_in_top5 from baseline_eval output."""
	per_ep = eval_data.get("per_episode", [])
	tier_buckets: dict[int, list[dict]] = {}
	for ep in per_ep:
	d = ep.get("difficulty", 1)
	tier_buckets.setdefault(d, []).append(ep)

	result = {}
	for tier, eps in tier_buckets.items():
	ndcgs = [e.get("ndcg", 0.0) for e in eps]
	fakes = [e.get("fakes_in_top5", 0) for e in eps]
	result[tier] = {
	"mean_ndcg": sum(ndcgs) / len(ndcgs),
	"mean_fakes": sum(fakes) / len(fakes),
	"n": len(eps),
	}
	return result


	def plot(
	training_log: Optional[str] = None,
	baseline_path: Optional[str] = None,
	trained_path: Optional[str] = None,
	output: str = "plots/results.png",
	wandb_url: Optional[str] = None,
	) -> None:
	import matplotlib
	matplotlib.use("Agg")
	import matplotlib.pyplot as plt
	import matplotlib.ticker as ticker
	import numpy as np

	Path(output).parent.mkdir(parents=True, exist_ok=True)

	# ── Determine number of subplots ──────────────────────────────────────────
	has_training = training_log is not None and Path(training_log).exists()
	has_before_after = (
	baseline_path is not None and Path(baseline_path).exists()
	and trained_path is not None and Path(trained_path).exists()
	)

	n_plots = 1 + (2 if has_before_after else 0)
	fig, axes = plt.subplots(n_plots, 1, figsize=(10, 4 * n_plots))
	if n_plots == 1:
	axes = [axes]

	fig.suptitle("Recruiter Agent Training Results", fontsize=14, fontweight="bold", y=0.98)

	# ── Subplot 1: Training curve ─────────────────────────────────────────────
	ax = axes[0]
	ax.set_title("Episode Return During Training")
	ax.set_xlabel("Episode")
	ax.set_ylabel("Return (reward)")

	if has_training:
	records = _load_jsonl(training_log)
	episodes = [r.get("episode", i * 4) for i, r in enumerate(records)]
	# Support both "mean_return" (train_grpo.py key) and "mean_reward" (legacy)
	rewards = [r.get("mean_terminal_return", r.get("mean_return", r.get("mean_reward", 0.0))) for r in records]
	smooth = _rolling_mean(rewards, window=10)

	ax.plot(episodes, rewards, alpha=0.3, color="steelblue", linewidth=0.8, label="raw")
	ax.plot(episodes, smooth, color="steelblue", linewidth=2.0, label="smoothed (win=10)")

	# Curriculum tier markers
	difficulties = [r.get("difficulty", 1) for r in records]
	prev_d = difficulties[0] if difficulties else 1
	for i, d in enumerate(difficulties):
	if d != prev_d:
	ax.axvline(x=episodes[i], color="orange", linestyle="--", alpha=0.6, linewidth=1.2)
	ax.text(episodes[i], ax.get_ylim()[0] if ax.get_ylim()[0] != 0 else min(rewards),
	f"T{d}", fontsize=8, color="darkorange", ha="left")
	prev_d = d
	else:
	ax.text(0.5, 0.5, "No training log provided", transform=ax.transAxes,
	ha="center", va="center", fontsize=12, color="grey")

	# Baseline reference line
	if baseline_path and Path(baseline_path).exists():
	baseline = _load_json(baseline_path)
	mean_base = baseline.get("summary", {}).get("mean_reward", None)
	if mean_base is not None:
	ax.axhline(y=mean_base, color="red", linestyle="--", linewidth=1.5,
	label=f"random baseline ({mean_base:.2f})")

	ax.legend(loc="lower right", fontsize=9)
	ax.grid(True, alpha=0.3)

	if has_before_after:
	baseline_data = _load_json(baseline_path)
	trained_data = _load_json(trained_path)
	before_stats = _tier_stats(baseline_data)
	after_stats = _tier_stats(trained_data)

	all_tiers = sorted(set(before_stats) \| set(after_stats))
	tier_labels = [f"T{t}" for t in all_tiers]
	x = np.arange(len(all_tiers))
	width = 0.35

	# ── Subplot 2: nDCG@5 per tier ────────────────────────────────────────
	ax2 = axes[1]
	ax2.set_title("nDCG@5 Before vs After Training (per difficulty tier)")
	ax2.set_ylabel("nDCG@5")
	ax2.set_ylim(0, 1.05)

	before_ndcg = [before_stats.get(t, {}).get("mean_ndcg", 0.0) for t in all_tiers]
	after_ndcg = [after_stats.get(t, {}).get("mean_ndcg", 0.0) for t in all_tiers]

	ax2.bar(x - width / 2, before_ndcg, width, label="Before (zero-shot)", color="salmon", alpha=0.8)
	ax2.bar(x + width / 2, after_ndcg, width, label="After (fine-tuned)", color="steelblue", alpha=0.8)
	ax2.set_xticks(x)
	ax2.set_xticklabels(tier_labels)
	ax2.legend(fontsize=9)
	ax2.grid(True, alpha=0.3, axis="y")

	# Annotate delta
	for i, (b, a) in enumerate(zip(before_ndcg, after_ndcg)):
	delta = a - b
	color = "green" if delta >= 0 else "red"
	ax2.text(x[i] + width / 2, a + 0.02, f"{delta:+.2f}", ha="center",
	fontsize=8, color=color, fontweight="bold")

	# ── Subplot 3: Fakes-in-top-5 per tier ───────────────────────────────
	ax3 = axes[2]
	ax3.set_title("Fakes in Top-5 Before vs After Training (per difficulty tier)")
	ax3.set_ylabel("Mean fakes in top-5 (lower is better)")

	before_fakes = [before_stats.get(t, {}).get("mean_fakes", 0.0) for t in all_tiers]
	after_fakes = [after_stats.get(t, {}).get("mean_fakes", 0.0) for t in all_tiers]

	ax3.bar(x - width / 2, before_fakes, width, label="Before (zero-shot)", color="salmon", alpha=0.8)
	ax3.bar(x + width / 2, after_fakes, width, label="After (fine-tuned)", color="steelblue", alpha=0.8)
	ax3.set_xticks(x)
	ax3.set_xticklabels(tier_labels)
	ax3.legend(fontsize=9)
	ax3.grid(True, alpha=0.3, axis="y")

	# Annotate delta
	for i, (b, a) in enumerate(zip(before_fakes, after_fakes)):
	delta = a - b
	color = "green" if delta <= 0 else "red"
	ax3.text(x[i] + width / 2, a + 0.02, f"{delta:+.2f}", ha="center",
	fontsize=8, color=color, fontweight="bold")

	if wandb_url:
	fig.text(0.5, 0.01, f"W&B run: {wandb_url}", ha="center", fontsize=8, color="grey")

	plt.tight_layout(rect=[0, 0.02, 1, 0.97])
	plt.savefig(output, dpi=150, bbox_inches="tight")
	print(f"Saved plot to {output}")


	def main() -> None:
	parser = argparse.ArgumentParser(description="Plot recruiter training results")
	parser.add_argument("--training", default=None, help="Path to training_log.jsonl")
	parser.add_argument("--baseline", default=None, help="Path to baseline_eval output JSON")
	parser.add_argument("--trained", default=None, help="Path to post-training eval JSON")
	parser.add_argument("--output", default="plots/results.png", help="Output PNG path")
	parser.add_argument("--wandb-url", default=None, help="W&B run URL to embed in figure")
	args = parser.parse_args()

	plot(
	training_log=args.training,
	baseline_path=args.baseline,
	trained_path=args.trained,
	output=args.output,
	wandb_url=args.wandb_url,
	)


	if __name__ == "__main__":
	main()