Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -202,6 +202,11 @@ _TOOL_CALL_RE = re.compile(
     r"<tool_call>\s*(\{.*?\})\s*</tool_call>",
     re.DOTALL,
 )
+_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL)
+_BARE_JSON_RE = re.compile(
+    r"\{\s*\"name\"\s*:\s*\"[^\"]+\"\s*,\s*\"arguments\"\s*:\s*\{.*?\}\s*\}",
+    re.DOTALL,
+)
 
 
 def _env_tool_callables(env: Any) -> list[Callable[..., str]]:
@@ -214,22 +219,61 @@ def _env_tool_callables(env: Any) -> list[Callable[..., str]]:
     return [getattr(env, name) for name in names]
 
 
+def _parse_one_tool_obj(obj: dict[str, Any]) -> dict[str, Any] | None:
+    name = obj.get("name") or obj.get("function", {}).get("name")
+    args = (
+        obj.get("arguments")
+        or obj.get("parameters")
+        or obj.get("function", {}).get("arguments")
+        or {}
+    )
+    if isinstance(args, str):
+        try:
+            args = json.loads(args)
+        except json.JSONDecodeError:
+            args = {"input": args}
+    if not name:
+        return None
+    return {"name": str(name), "arguments": args if isinstance(args, dict) else {}}
+
+
 def _parse_tool_calls(text: str) -> list[dict[str, Any]]:
+    """Parse Qwen-style tool calls: <tool_call>, ```json fences, or bare {name,arguments}."""
     calls: list[dict[str, Any]] = []
-    for m in _TOOL_CALL_RE.finditer(text or ""):
+    seen: set[str] = set()
+    raw = text or ""
+
+    candidates: list[str] = []
+    candidates.extend(m.group(1) for m in _TOOL_CALL_RE.finditer(raw))
+    candidates.extend(m.group(1) for m in _JSON_FENCE_RE.finditer(raw))
+    candidates.extend(m.group(0) for m in _BARE_JSON_RE.finditer(raw))
+
+    # Whole-string JSON object / list fallback
+    stripped = raw.strip()
+    if stripped.startswith("{") or stripped.startswith("["):
+        candidates.append(stripped)
+
+    for blob in candidates:
         try:
-            obj = json.loads(m.group(1))
-            name = obj.get("name") or obj.get("function", {}).get("name")
-            args = obj.get("arguments") or obj.get("parameters") or obj.get("function", {}).get("arguments") or {}
-            if isinstance(args, str):
-                try:
-                    args = json.loads(args)
-                except json.JSONDecodeError:
-                    args = {"input": args}
-            if name:
-                calls.append({"name": str(name), "arguments": args if isinstance(args, dict) else {}})
+            obj = json.loads(blob)
         except json.JSONDecodeError:
             continue
+        items = obj if isinstance(obj, list) else [obj]
+        for item in items:
+            if not isinstance(item, dict):
+                continue
+            # unwrap {"type":"function","function":{...}}
+            if "function" in item and isinstance(item["function"], dict):
+                parsed = _parse_one_tool_obj(item["function"])
+            else:
+                parsed = _parse_one_tool_obj(item)
+            if parsed is None:
+                continue
+            key = json.dumps(parsed, sort_keys=True)
+            if key in seen:
+                continue
+            seen.add(key)
+            calls.append(parsed)
     return calls
 
 
@@ -994,8 +1038,8 @@ def main(cfg: Cfg | None = None) -> None:
         bf16=use_bf16,
         fp16=use_fp16,
         remove_unused_columns=False,
-        # log_completions_reward (0) + DeepAgentSWEEnv.get_reward (1)
-        reward_weights=[0.0, 1.0],
+        # Do not set reward_weights here: env get_reward is appended after init,
+        # so a fixed [0,1] length mismatches and crashes GRPOTrainer.
         hub_model_id=cfg.hub_model_id,
         hub_token=os.environ["HF_TOKEN"],
         push_to_hub=False,