Genesis-2.0 RLHF Phase 0 artifacts

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+dpo_pairs_all.jsonl filter=lfs diff=lfs merge=lfs -text
+dpo_pairs_perturbed.jsonl filter=lfs diff=lfs merge=lfs -text
+dpo_pairs_scored.jsonl filter=lfs diff=lfs merge=lfs -text

build_pairs.py

@@ -0,0 +1,307 @@
+"""
+Genesis-2.0 RLHF — Preference Pair Builder
+
+Generates DPO training pairs (chosen/rejected) from existing SFT data.
+
+Strategy:
+1. Use SFT prompts as seeds
+2. Generate multiple responses from Genesis-1.0 (via MLX on Mac or via API)
+3. Score with rule-based reward functions
+4. Best = chosen, worst = rejected
+
+For Phase 0 on MacBook, we use an offline approach:
+- Source A: Direct from SFT data (the existing trajectory IS the chosen)
+  Generate a perturbed version as rejected
+- Source B: Score-based (take existing trajectories, rank by reward score,
+  pair high/low within each prompt group)
+"""
+
+import json
+import os
+import random
+import sys
+from typing import Optional
+
+# Add project to path
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from rewards import combined_reward, reward_debug, extract_tool_calls
+
+
+DATA_DIR = "/Volumes/this_and_that/hermes-admin/improvements/hermes-agentic-dataset/data/train"
+OUTPUT_DIR = "/Users/jacobeen/model-forge-workspace/genesis-rlhf"
+
+
+def load_sft_data(sources: Optional[list[str]] = None) -> list[dict]:
+    """
+    Load SFT data from JSONL files.
+    Each entry: {"text": "...", "metadata": {...}}
+    """
+    if sources is None:
+        sources = [
+            "train_sessions_00001.jsonl",
+            "train_augmented_00001.jsonl",
+        ]
+
+    data = []
+    for src in sources:
+        path = os.path.join(DATA_DIR, src)
+        if not os.path.exists(path):
+            print(f"  WARNING: {path} not found, skipping")
+            continue
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    data.append(json.loads(line))
+    return data
+
+
+def extract_prompt(text: str) -> str:
+    """Extract the user prompt from a full conversation text.
+    Returns everything up to but NOT including the first <|im_start|>assistant."""
+    idx = text.find("<|im_start|>assistant")
+    if idx >= 0:
+        return text[:idx].strip()
+    return text
+
+
+def extract_completion(text: str) -> str:
+    """Extract the assistant's completion, including the assistant header."""
+    idx = text.find("<|im_start|>assistant")
+    if idx >= 0:
+        return text[idx:].strip()
+    return text
+
+
+def perturb_trajectory(text: str) -> str:
+    """
+    Create a deliberately worse version of a trajectory for DPO rejected pairs.
+    Uses aggressive perturbations that significantly degrade quality.
+    """
+    import re
+
+    # Use multiple perturbations together for stronger signal
+    strategies = []
+
+    # Always remove some thinking
+    text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
+
+    tool_blocks = list(re.finditer(r'<tool_call>.*?</tool_call>', text, re.DOTALL))
+
+    if tool_blocks:
+        # 50% chance: remove the last tool call entirely (drops answer)
+        if random.random() < 0.5 and len(tool_blocks) >= 1:
+            idx = len(tool_blocks) - 1
+            start, end = tool_blocks[idx].start(), tool_blocks[idx].end()
+            text = text[:start] + text[end:]
+
+        # 40% chance: replace a tool name with a plausible wrong one
+        if random.random() < 0.4:
+            tool_names = ["web_search", "web_extract", "code_interpreter", "file_read",
+                           "file_write", "database_query", "send_email", "calculator",
+                           "search_web", "fetch_url", "run_code", "read_document"]
+            def _replace_name(m):
+                try:
+                    tc = json.loads(m.group(1))
+                    others = [n for n in tool_names if n != tc.get("name", "")]
+                    if others:
+                        tc["name"] = random.choice(others)
+                    return f'<tool_call>\n{json.dumps(tc, indent=2)}\n</tool_call>'
+                except:
+                    return m.group(0)
+            text = re.sub(r'<tool_call>\s*(\{.*?\})\s*</tool_call>', _replace_name, text, count=1, flags=re.DOTALL)
+
+        # 40% chance: make arguments invalid
+        if random.random() < 0.4:
+            def _corrupt_args(m):
+                try:
+                    tc = json.loads(m.group(1))
+                    if "arguments" in tc and isinstance(tc["arguments"], dict):
+                        # Remove a required-looking argument
+                        for key in list(tc["arguments"].keys())[:1]:
+                            del tc["arguments"][key]
+                            break
+                    return f'<tool_call>\n{json.dumps(tc, indent=2)}\n</tool_call>'
+                except:
+                    return m.group(0)
+            text = re.sub(r'<tool_call>\s*(\{.*?\})\s*</tool_call>', _corrupt_args, text, count=1, flags=re.DOTALL)
+
+    # Remove answer text after the last tool call
+    if random.random() < 0.5:
+        parts = text.rsplit("<|im_start|>assistant\n", 1)
+        if len(parts) > 1:
+            last_asst = parts[1]
+            tc_end = last_asst.rfind("</tool_call>")
+            if tc_end >= 0:
+                last_asst = last_asst[:tc_end + len("</tool_call>")]
+                text = parts[0] + "<|im_start|>assistant\n" + last_asst
+            else:
+                # No tool call at all — just drop the answer
+                text = parts[0].strip()
+
+    text = re.sub(r'\n{3,}', '\n\n', text)
+    return text
+
+
+def build_pairs_scored(
+    data: list[dict],
+    output_path: str,
+    score_threshold: float = 0.3,
+    max_pairs: int = 2000,
+) -> list[dict]:
+    """
+    Build DPO pairs by scoring existing trajectories and pairing
+    high-scoring vs low-scoring examples.
+
+    Each pair: {"prompt": ..., "chosen": ..., "rejected": ...}
+    """
+    scored = []
+    for item in data:
+        score = combined_reward(item["text"])
+        prompt = extract_prompt(item["text"])
+        completion = extract_completion(item["text"])
+        scored.append((score, prompt, completion, item["metadata"]))
+
+    # Sort by score
+    scored.sort(key=lambda x: x[0], reverse=True)
+
+    pairs = []
+    # Pair high with low
+    high_idx = 0
+    low_idx = len(scored) - 1
+
+    while high_idx < low_idx and len(pairs) < max_pairs:
+        high_score, high_prompt, high_comp, high_meta = scored[high_idx]
+        low_score, low_prompt, low_comp, low_meta = scored[low_idx]
+
+        score_gap = high_score - low_score
+        if score_gap >= score_threshold and high_score > 0.5 and low_score < 0.8:
+            pair = {
+                "prompt": high_prompt,
+                "chosen": high_comp,
+                "rejected": low_comp,
+                "score_chosen": high_score,
+                "score_rejected": low_score,
+                "metadata": {
+                    "source_high": high_meta.get("source", ""),
+                    "source_low": low_meta.get("source", ""),
+                }
+            }
+            pairs.append(pair)
+
+        high_idx += 1
+        low_idx -= 1
+
+    # Save
+    with open(output_path, "w") as f:
+        for p in pairs:
+            f.write(json.dumps(p) + "\n")
+
+    print(f"Built {len(pairs)} scored pairs → {output_path}")
+    print(f"  Score range: {scored[0][0]:.3f} (high) to {scored[-1][0]:.3f} (low)")
+    return pairs
+
+
+def build_pairs_perturbed(
+    data: list[dict],
+    output_path: str,
+    max_pairs: int = 2000,
+) -> list[dict]:
+    """
+    Build DPO pairs by taking existing trajectories and creating
+    perturbed (deliberately worse) versions as the rejected sample.
+    The original trajectory is the chosen sample.
+    """
+    pairs = []
+    for item in data:
+        if len(pairs) >= max_pairs:
+            break
+
+        prompt = extract_prompt(item["text"])
+        chosen = extract_completion(item["text"])
+
+        # Verify chosen scores well (skip bad data)
+        chosen_score = combined_reward(chosen)
+        if chosen_score < 0.5:
+            continue
+
+        # Create perturbed version
+        rejected = perturb_trajectory(chosen)
+        rejected_score = combined_reward(rejected)
+
+        # Only accept if the perturbation actually made it worse
+        if rejected_score < chosen_score - 0.05:
+            pair = {
+                "prompt": prompt,
+                "chosen": chosen,
+                "rejected": rejected,
+                "score_chosen": chosen_score,
+                "score_rejected": rejected_score,
+                "metadata": {
+                    "source": item["metadata"].get("source", "perturbed"),
+                }
+            }
+            pairs.append(pair)
+
+    with open(output_path, "w") as f:
+        for p in pairs:
+            f.write(json.dumps(p) + "\n")
+
+    print(f"Built {len(pairs)} perturbed pairs → {output_path}")
+    return pairs
+
+
+if __name__ == "__main__":
+    print("=== Genesis-2.0: DPO Preference Pair Builder ===\n")
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    # Load data
+    print("Loading SFT data...")
+    data = load_sft_data()
+
+    # Sample sharegpt-fc (too large to load all — sample 500)
+    sharegpt_path = os.path.join(DATA_DIR, "train_hf_sharegpt-fc.jsonl")
+    if os.path.exists(sharegpt_path):
+        with open(sharegpt_path) as f:
+            for i, line in enumerate(f):
+                if i >= 500:
+                    break
+                data.append(json.loads(line))
+        print(f"  + 500 from sharegpt-fc")
+
+    print(f"  Loaded {len(data)} total examples\n")
+
+    # Method 1: Score-based pairing
+    print("Building scored pairs...")
+    build_pairs_scored(
+        data,
+        os.path.join(OUTPUT_DIR, "dpo_pairs_scored.jsonl"),
+        max_pairs=1000,
+    )
+
+    # Method 2: Perturbation-based pairing
+    print("\nBuilding perturbed pairs...")
+    build_pairs_perturbed(
+        data,
+        os.path.join(OUTPUT_DIR, "dpo_pairs_perturbed.jsonl"),
+        max_pairs=1000,
+    )
+
+    # Combine
+    print("\nCombining...")
+    combined = []
+    for method in ["scored", "perturbed"]:
+        path = os.path.join(OUTPUT_DIR, f"dpo_pairs_{method}.jsonl")
+        if os.path.exists(path):
+            with open(path) as f:
+                for line in f:
+                    if line.strip():
+                        combined.append(json.loads(line))
+
+    with open(os.path.join(OUTPUT_DIR, "dpo_pairs_all.jsonl"), "w") as f:
+        for p in combined:
+            f.write(json.dumps(p) + "\n")
+
+    print(f"\n=== DONE: {len(combined)} total DPO pairs ===")
+    print(f"  File: {OUTPUT_DIR}/dpo_pairs_all.jsonl")

dpo_pairs_all.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cf41af65be6b23b77a533dc70a367ce24a2197b7bf0e12d9a8f68638db912a4
+size 43781476

dpo_pairs_perturbed.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:832766423dc534ebbee155d47370c1af57c73a773cc2f422be9b3a67d1b6d934
+size 26361675

dpo_pairs_scored.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a5850f0aa43e90895f1f2354348514c9a1a6279704eb9e56e8070e07b7e8abd4
+size 17419801

eval_benchmark.py

@@ -0,0 +1,199 @@
+"""
+Genesis-2.0 RLHF — Evaluation Benchmark
+
+Evaluates a model's tool-use capabilities on a hold-out set.
+Measures: tool accuracy, completion rate, efficiency, format compliance.
+
+Usage:
+    python3 eval_benchmark.py --model <path> [--data <jsonl>] [--output <json>]
+"""
+
+import json
+import os
+import sys
+import re
+from typing import Optional
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from rewards import (
+    extract_tool_calls,
+    extract_assistant_messages,
+    combined_reward,
+    reward_debug,
+)
+
+
+def evaluate_trajectory(trajectory_text: str) -> dict:
+    """
+    Evaluate a single model trajectory on all metrics.
+    Returns a dict of scores and metadata.
+    """
+    tool_calls = extract_tool_calls(trajectory_text)
+    assistant_msgs = extract_assistant_messages(trajectory_text)
+    reward_scores = reward_debug(trajectory_text)
+
+    # Tool accuracy: how many tool calls have valid names + args
+    valid_calls = 0
+    for tc in tool_calls:
+        name_ok = isinstance(tc.get("name"), str) and bool(tc["name"].strip())
+        args_ok = "arguments" in tc and isinstance(tc["arguments"], dict)
+        if name_ok and args_ok:
+            valid_calls += 1
+
+    tool_accuracy = valid_calls / max(len(tool_calls), 1)
+
+    # Completion: is there a substantive answer?
+    if assistant_msgs:
+        last_msg = assistant_msgs[-1]
+        has_text_answer = len(re.sub(r'<[^>]+>', '', last_msg).strip()) > 30
+    else:
+        has_text_answer = False
+
+    return {
+        "tool_count": len(tool_calls),
+        "valid_tool_count": valid_calls,
+        "tool_accuracy": tool_accuracy,
+        "has_answer": 1.0 if has_text_answer else 0.0,
+        "reward_completion": reward_scores["completion"],
+        "reward_validity": reward_scores["validity"],
+        "reward_efficiency": reward_scores["efficiency"],
+        "reward_format": reward_scores["format"],
+        "reward_combined": reward_scores["combined"],
+        "assistant_turns": len(assistant_msgs),
+    }
+
+
+def evaluate_dataset(
+    data_path: str,
+    model_generate_fn=None,
+    output_path: Optional[str] = None,
+    max_samples: int = 200,
+) -> dict:
+    """
+    Evaluate a dataset of prompts against a model.
+
+    If model_generate_fn is None, evaluates the existing trajectories
+    (for offline eval of SFT data).
+
+    If model_generate_fn is provided, it's called as:
+        response = model_generate_fn(prompt)
+    for each prompt.
+
+    Returns aggregated metrics.
+    """
+    # Load data
+    data = []
+    with open(data_path) as f:
+        for i, line in enumerate(f):
+            if i >= max_samples:
+                break
+            line = line.strip()
+            if line:
+                data.append(json.loads(line))
+
+    results = []
+    for item in data:
+        if model_generate_fn:
+            prompt = item.get("prompt", item.get("text", ""))
+            response = model_generate_fn(prompt)
+            full_text = prompt + "\n" + response
+        else:
+            full_text = item.get("text", "")
+
+        metrics = evaluate_trajectory(full_text)
+        results.append(metrics)
+
+    # Aggregate
+    agg = {
+        "num_samples": len(results),
+        "avg_tool_accuracy": sum(r["tool_accuracy"] for r in results) / max(len(results), 1),
+        "avg_has_answer": sum(r["has_answer"] for r in results) / max(len(results), 1),
+        "avg_reward_combined": sum(r["reward_combined"] for r in results) / max(len(results), 1),
+        "avg_reward_completion": sum(r["reward_completion"] for r in results) / max(len(results), 1),
+        "avg_reward_validity": sum(r["reward_validity"] for r in results) / max(len(results), 1),
+        "avg_reward_efficiency": sum(r["reward_efficiency"] for r in results) / max(len(results), 1),
+        "avg_reward_format": sum(r["reward_format"] for r in results) / max(len(results), 1),
+        "avg_tool_count": sum(r["tool_count"] for r in results) / max(len(results), 1),
+        "total_tool_calls": sum(r["tool_count"] for r in results),
+    }
+
+    if output_path:
+        with open(output_path, "w") as f:
+            json.dump({"aggregate": agg, "per_sample": results}, f, indent=2)
+        print(f"Saved evaluation to {output_path}")
+
+    return agg
+
+
+def build_holdout_set(
+    input_dir: str,
+    output_path: str,
+    num_prompts: int = 150,
+    sources: Optional[list[str]] = None,
+) -> None:
+    """
+    Build a hold-out evaluation set from SFT data.
+    Takes prompts only (no completions) for online LLM evaluation.
+    """
+    if sources is None:
+        sources = ["train_sessions_00001.jsonl", "train_augmented_00001.jsonl"]
+
+    prompts = []
+    for src in sources:
+        path = os.path.join(input_dir, src)
+        if not os.path.exists(path):
+            print(f"  WARNING: {path} not found, skipping")
+            continue
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if line:
+                    item = json.loads(line)
+                    # Extract just the prompt (user + system, no assistant)
+                    text = item["text"]
+                    idx = text.find("<|im_start|>assistant")
+                    if idx >= 0:
+                        prompt = text[:idx].strip()
+                    else:
+                        prompt = text
+                    prompts.append({
+                        "prompt": prompt,
+                        "source": item["metadata"].get("source", "unknown"),
+                    })
+
+    # Take evenly spaced samples
+    step = max(1, len(prompts) // num_prompts)
+    holdout = [prompts[i] for i in range(0, len(prompts), step)][:num_prompts]
+
+    with open(output_path, "w") as f:
+        for p in holdout:
+            f.write(json.dumps(p) + "\n")
+
+    print(f"Built hold-out set: {len(holdout)} prompts → {output_path}")
+
+
+if __name__ == "__main__":
+    project_dir = os.path.dirname(os.path.abspath(__file__))
+    data_dir = "/Volumes/this_and_that/hermes-admin/improvements/hermes-agentic-dataset/data/train"
+
+    print("=== Genesis-2.0: Evaluation Benchmark ===\n")
+
+    # Build hold-out set
+    holdout_path = os.path.join(project_dir, "eval_holdout.jsonl")
+    build_holdout_set(data_dir, holdout_path, num_prompts=150)
+    print()
+
+    # Evaluate SFT data offline (baseline)
+    for src in ["train_sessions_00001.jsonl", "train_augmented_00001.jsonl"]:
+        path = os.path.join(data_dir, src)
+        if os.path.exists(path):
+            print(f"\nEvaluating {src}...")
+            agg = evaluate_dataset(path, max_samples=50)
+            print(f"  Tool accuracy: {agg['avg_tool_accuracy']:.3f}")
+            print(f"  Completion rate: {agg['avg_has_answer']:.3f}")
+            print(f"  Combined reward: {agg['avg_reward_combined']:.3f}")
+            print(f"  Tool efficiency: {agg['avg_reward_efficiency']:.3f}")
+            print(f"  Format compliance: {agg['avg_reward_format']:.3f}")
+
+    print(f"\nHold-out set ready: {holdout_path}")
+    print(f"Run evaluation with: python3 eval_benchmark.py --eval {holdout_path}")

rewards.py

@@ -0,0 +1,221 @@
+"""
+Genesis-2.0 RLHF — Reward Functions Module
+
+Rule-based reward functions for scoring tool-use trajectories.
+Used by both the DPO preference pair generator and the GRPO trainer.
+"""
+
+import json
+import re
+from typing import Any, Optional
+
+
+def extract_tool_calls(text: str) -> list[dict[str, Any]]:
+    """Extract <tool_call> JSON blocks from assistant messages."""
+    tool_calls = []
+    pattern = re.compile(r'<tool_call>\s*(\{.*?\})\s*</tool_call>', re.DOTALL)
+    for match in pattern.finditer(text):
+        try:
+            tc = json.loads(match.group(1))
+            if isinstance(tc, dict) and "name" in tc:
+                tool_calls.append(tc)
+        except json.JSONDecodeError:
+            continue
+    return tool_calls
+
+
+def extract_assistant_messages(text: str) -> list[str]:
+    """Split a conversation into individual assistant response blocks."""
+    parts = text.split("<|im_start|>assistant\n")
+    if len(parts) <= 1:
+        return []
+    messages = []
+    for p in parts[1:]:
+        end_idx = p.find("<|im_end|>")
+        if end_idx >= 0:
+            messages.append(p[:end_idx].strip())
+    return messages
+
+
+def reward_tool_call_validity(trajectory: str) -> float:
+    """
+    +1.0 if all <tool_call> blocks contain valid JSON with valid tool names.
+    0.0 if ANY tool call has invalid JSON or missing 'name' field.
+    """
+    tool_calls = extract_tool_calls(trajectory)
+    if not tool_calls:
+        return 0.0
+    for tc in tool_calls:
+        if not isinstance(tc.get("name"), str) or not tc["name"].strip():
+            return 0.0
+        if "arguments" not in tc:
+            return 0.5  # partial: has name but no args
+    return 1.0
+
+
+def reward_task_completion(trajectory: str, expected_keywords: Optional[list[str]] = None) -> float:
+    """
+    +1.0 if the final assistant message provides a substantive answer.
+    0.0 if the trajectory ends without a meaningful response.
+    -1.0 if it ends with an error or hallucination.
+    """
+    msgs = extract_assistant_messages(trajectory)
+    if not msgs:
+        return -1.0
+
+    last = msgs[-1]
+
+    # If the last message has tool calls but no text answer, task isn't complete
+    text_after_tools = re.sub(r'<tool_call>.*?</tool_call>', '', last, flags=re.DOTALL).strip()
+    text_after_tools = re.sub(r'<tool_result>.*?</tool_result>', '', text_after_tools, flags=re.DOTALL).strip()
+
+    tool_calls = extract_tool_calls(last)
+    if not text_after_tools and tool_calls:
+        return 0.0  # only tool calls, no answer
+
+    if text_after_tools and len(text_after_tools) > 20:
+        # Has substantive text — likely gave an answer
+        if expected_keywords:
+            hits = sum(1 for kw in expected_keywords if kw.lower() in text_after_tools.lower())
+            return min(1.0, hits / max(len(expected_keywords), 1) * 1.5)
+        return 1.0
+
+    if text_after_tools and len(text_after_tools) > 5:
+        return 0.5  # short answer, partial completion
+
+    return 0.0
+
+
+def reward_tool_result_usage(trajectory: str) -> float:
+    """
+    +1.0 if tool calls are followed by reasoning that references them.
+    0.0 if tool calls are made but no subsequent reasoning.
+    """
+    segments = trajectory.split("<|im_start|>assistant\n")
+    tool_call_segments = 0
+    reasoning_after_call = 0
+
+    for seg in segments[1:]:
+        has_tool_call = bool(re.search(r'<tool_call>', seg))
+        if has_tool_call:
+            tool_call_segments += 1
+            # Remove tool_call blocks and check remaining text
+            text_after = re.sub(r'<tool_call>.*?</tool_call>', '', seg, flags=re.DOTALL)
+            text_after = re.sub(r'<think>.*?</think>', '', text_after, flags=re.DOTALL)
+            text_after = text_after.strip()
+            if len(text_after) > 50:
+                reasoning_after_call += 1
+
+    if tool_call_segments == 0:
+        return 0.5  # no tools used — neutral
+    if reasoning_after_call >= tool_call_segments * 0.5:
+        return 1.0
+    if reasoning_after_call > 0:
+        return 0.5
+    return 0.0
+
+
+def reward_trajectory_efficiency(trajectory: str) -> float:
+    """
+    Score based on how efficiently tools are used.
+    Ideal: 1-3 tool calls for a typical query.
+    Penalty for excessive tool calls (>5) or zero tool calls when tools would help.
+    """
+    tool_calls = extract_tool_calls(trajectory)
+    n = len(tool_calls)
+
+    if n == 0:
+        return 0.3  # no tools used — might be fine for simple queries
+    elif n <= 3:
+        return 1.0  # efficient
+    elif n <= 5:
+        return 0.7  # acceptable
+    elif n <= 8:
+        return 0.4  # somewhat excessive
+    else:
+        return 0.0  # too many tool calls
+
+
+def reward_format_validity(trajectory: str) -> float:
+    """
+    +1.0 if tool calls use the correct <tool_call> JSON format.
+    0.0 if there are malformed tool calls.
+    """
+    # Check for malformed tool_call markers
+    malformed = re.findall(r'<tool_call>[^<]*?$', trajectory)
+    if malformed:
+        return 0.0
+
+    # Check for unclosed tool_call tags
+    open_tags = len(re.findall(r'<tool_call>', trajectory))
+    close_tags = len(re.findall(r'</tool_call>', trajectory))
+    if open_tags != close_tags:
+        return 0.0
+
+    # All parsed tool calls should have valid JSON
+    tool_calls = extract_tool_calls(trajectory)
+    expected = len(re.findall(r'<tool_call>', trajectory))
+    if expected > 0 and len(tool_calls) < expected:
+        return 0.5  # some valid, some not
+
+    return 1.0
+
+
+def reward_recovery_from_error(trajectory: str) -> float:
+    """
+    +1.0 if the model retries after a failed tool call (multi-turn recovery).
+    Detects consecutive assistant segments making tool calls.
+    0.0 otherwise.
+    """
+    segments = trajectory.split("<|im_start|>assistant\n")[1:]
+    consecutive_tool_calls = 0
+    for i, seg in enumerate(segments):
+        has_tc = bool(re.search(r'<tool_call>', seg))
+        if has_tc and i > 0:
+            prev_has_tc = bool(re.search(r'<tool_call>', segments[i-1]))
+            if prev_has_tc:
+                consecutive_tool_calls += 1
+
+    if consecutive_tool_calls >= 1:
+        return 1.0  # did a multi-step tool sequence
+    return 0.0
+
+
+def combined_reward(trajectory: str, expected_keywords: Optional[list[str]] = None) -> float:
+    """
+    Weighted combination of all reward signals.
+    Used by both DPO pair generation and GRPO training.
+    """
+    weights = {
+        "validity": 0.25,
+        "completion": 0.40,
+        "usage": 0.15,
+        "efficiency": 0.10,
+        "format": 0.05,
+        "recovery": 0.05,
+    }
+
+    scores = {
+        "validity": reward_tool_call_validity(trajectory),
+        "completion": reward_task_completion(trajectory, expected_keywords),
+        "usage": reward_tool_result_usage(trajectory),
+        "efficiency": reward_trajectory_efficiency(trajectory),
+        "format": reward_format_validity(trajectory),
+        "recovery": reward_recovery_from_error(trajectory),
+    }
+
+    total = sum(scores[k] * weights[k] for k in weights)
+    return max(-1.0, min(1.0, total))  # clamp to [-1, 1]
+
+
+def reward_debug(trajectory: str) -> dict:
+    """Return breakdown of all reward components for debugging."""
+    return {
+        "validity": reward_tool_call_validity(trajectory),
+        "completion": reward_task_completion(trajectory),
+        "usage": reward_tool_result_usage(trajectory),
+        "efficiency": reward_trajectory_efficiency(trajectory),
+        "format": reward_format_validity(trajectory),
+        "recovery": reward_recovery_from_error(trajectory),
+        "combined": combined_reward(trajectory),
+    }

train_grpo.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+"""
+Genesis-2.0 — Phase 2: GRPO with Verifiable Rewards
+
+Run on RunPod RTX PRO 6000 Blackwell (96 GB) using TRL + Unsloth + vLLM.
+
+Steps:
+1. Load Phase 1 checkpoint (or Genesis-1.0 base + all-linear LoRA)
+2. Configure GRPOTrainer with vLLM for rollouts
+3. Run GRPO with rule-based reward functions
+4. Save + upload adapter
+
+Usage:
+    python3 train_grpo.py [--from-dpo /workspace/genesis2-dpo]
+"""
+
+import os
+import sys
+import json
+import torch
+import argparse
+from typing import Optional
+
+os.environ["UNSLOTH_VLLM_STANDBY"] = "1"
+
+
+class Config:
+    # Model
+    base_model = "Qwen/Qwen3.6-35B-A3B"
+    # If --from-dpo is set, load that adapter. Otherwise load Genesis-1.0 SFT adapter
+    initial_adapter = "jacobeen06/Genesis-1.0-SFT-adapter"
+    dpo_adapter_path = None  # Override with --from-dpo
+
+    # Output
+    output_dir = "/workspace/genesis2-grpo"
+    hf_repo = "jacobeen06/Genesis-2.0-GRPO-adapter"
+
+    # QLoRA
+    load_in_4bit = True
+    bnb_4bit_quant_type = "nf4"
+    bnb_4bit_compute_dtype = torch.bfloat16
+
+    # All-linear LoRA (Config C)
+    lora_r = 32
+    lora_alpha = 64
+    lora_dropout = 0.0
+    lora_target_modules = [
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
+        "gate",
+    ]
+    use_rslora = True
+
+    # GRPO
+    num_generations = 4  # G=4 responses per prompt
+    max_length = 4096
+    max_prompt_length = 3072
+    beta = 0.04  # KL penalty (start)
+    beta_decay = True  # Decay to 0.01 during training
+    clip_high = 0.28  # DAPO-style asymmetric
+    clip_low = 0.20
+
+    # Training
+    learning_rate = 3e-6
+    lr_scheduler_type = "cosine"
+    warmup_ratio = 0.05
+    per_device_train_batch_size = 1
+    gradient_accumulation_steps = 4
+    num_train_epochs = 1
+    logging_steps = 5
+    save_steps = 100
+    save_total_limit = 2
+
+    # vLLM
+    vllm_gpu_memory_utilization = 0.90
+
+    # Data
+    prompts_data = "/workspace/training_prompts.jsonl"  # Prompts for GRPO rollout
+
+
+def load_model(config: Config):
+    """Load base model + merge Genesis-1.0 SFT → attach all-linear LoRA."""
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+    from peft import PeftModel, LoraConfig, get_peft_model
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=config.load_in_4bit,
+        bnb_4bit_quant_type=config.bnb_4bit_quant_type,
+        bnb_4bit_compute_dtype=config.bnb_4bit_compute_dtype,
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        config.base_model,
+        quantization_config=bnb_config,
+        device_map="auto",
+        trust_remote_code=True,
+        torch_dtype=torch.bfloat16,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(config.base_model, trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Load adapter
+    if config.dpo_adapter_path and os.path.exists(config.dpo_adapter_path):
+        print(f"Loading DPO adapter from {config.dpo_adapter_path}")
+        model = PeftModel.from_pretrained(model, config.dpo_adapter_path)
+    else:
+        print(f"Loading Genesis-1.0 SFT adapter from {config.initial_adapter}")
+        model = PeftModel.from_pretrained(model, config.initial_adapter)
+
+    model = model.merge_and_unload()
+
+    # Attach new LoRA for GRPO
+    lora_config = LoraConfig(
+        r=config.lora_r,
+        lora_alpha=config.lora_alpha,
+        target_modules=config.lora_target_modules,
+        lora_dropout=config.lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+        use_rslora=config.use_rslora,
+    )
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    return model, tokenizer
+
+
+def reward_func(completions, **kwargs):
+    """
+    Rule-based reward function for GRPO.
+    Runs the rewards module on each generated completion.
+    """
+    # Import our reward functions
+    sys.path.insert(0, "/workspace/genesis-rlhf")
+    from rewards import combined_reward
+
+    rewards = []
+    for completion in completions:
+        score = combined_reward(completion)
+        rewards.append(score)
+    return rewards
+
+
+def load_prompts(config: Config):
+    """Load training prompts for GRPO rollout."""
+    prompts = []
+    with open(config.prompts_data) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                item = json.loads(line)
+                prompts.append(item.get("prompt", item.get("text", "")))
+    print(f"Loaded {len(prompts)} prompts for GRPO")
+    return prompts
+
+
+def train_grpo(config: Config):
+    from trl import GRPOTrainer
+    from transformers import TrainingArguments
+    from datasets import Dataset
+
+    model, tokenizer = load_model(config)
+    prompts = load_prompts(config)
+
+    # Create dataset with just prompts
+    dataset = Dataset.from_list([{"prompt": p} for p in prompts])
+
+    # Training args
+    training_args = TrainingArguments(
+        output_dir=config.output_dir,
+        per_device_train_batch_size=config.per_device_train_batch_size,
+        gradient_accumulation_steps=config.gradient_accumulation_steps,
+        learning_rate=config.learning_rate,
+        lr_scheduler_type=config.lr_scheduler_type,
+        warmup_ratio=config.warmup_ratio,
+        num_train_epochs=config.num_train_epochs,
+        logging_steps=config.logging_steps,
+        save_steps=config.save_steps,
+        save_total_limit=config.save_total_limit,
+        bf16=True,
+        tf32=True,
+        gradient_checkpointing=True,
+        gradient_checkpointing_kwargs={"use_reentrant": False},
+        report_to="wandb" if os.environ.get("WANDB_API_KEY") else "none",
+        run_name="genesis2-grpo",
+        dataloader_num_workers=2,
+    )
+
+    # GRPO Trainer
+    trainer = GRPOTrainer(
+        model=model,
+        reward_funcs=[reward_func],
+        args=training_args,
+        train_dataset=dataset,
+        tokenizer=tokenizer,
+        num_generations=config.num_generations,
+        max_length=config.max_length,
+        max_prompt_length=config.max_prompt_length,
+        beta=config.beta,
+        clip_high=config.clip_high,
+        clip_low=config.clip_low,
+    )
+
+    # Train
+    print("Starting GRPO training...")
+    trainer.train()
+
+    # Save
+    trainer.save_model(config.output_dir)
+    tokenizer.save_pretrained(config.output_dir)
+
+    # Upload
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.create_repo(config.hf_repo, exist_ok=True)
+        api.upload_folder(
+            folder_path=config.output_dir,
+            repo_id=config.hf_repo,
+            commit_message="Genesis-2.0 GRPO adapter",
+        )
+        print(f"Uploaded to {config.hf_repo}")
+    except Exception as e:
+        print(f"Upload failed (non-fatal): {e}")
+
+    print(f"DONE! Adapter saved to {config.output_dir}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--from-dpo", type=str, default=None,
+                        help="Path to DPO-trained adapter to start from")
+    args = parser.parse_args()
+
+    config = Config()
+    config.dpo_adapter_path = args.from_dpo
+
+    print("=" * 60)
+    print("Genesis-2.0 — Phase 2: GRPO Training")
+    print("=" * 60)
+    print(f"Base model: {config.base_model}")
+    print(f"Start from DPO: {config.dpo_adapter_path or 'No (from SFT)'}")
+    print(f"Generations per prompt: {config.num_generations}")
+    print(f"Output: {config.output_dir}")
+    print()
+
+    train_grpo(config)