crca 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

training/datasets.py

@@ -0,0 +1,140 @@
+"""Dataset assembly for ReAct training traces."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional, Sequence
+
+from crca_reasoning.types import LRMPlanResult
+
+
+@dataclass
+class ReActExample:
+    prompt: str
+    response: str
+    tags: Dict[str, str]
+    refusal: bool = False
+
+
+def plan_result_to_examples(plan: LRMPlanResult) -> List[ReActExample]:
+    examples: List[ReActExample] = []
+    for cycle in plan.cycle_traces:
+        prompt = cycle.reasoning
+        response = ""
+        if cycle.actions:
+            response += "Actions:\n"
+            for act in cycle.actions:
+                response += f"- {act.tool_name}: {act.payload}\n"
+        if cycle.critique:
+            response += f"Critique: {cycle.critique}\n"
+        examples.append(
+            ReActExample(
+                prompt=prompt,
+                response=response,
+                tags={"type": "react_cycle"},
+                refusal=any(obs.refusal is not None for obs in cycle.observations),
+            )
+        )
+    if plan.rationale_trace and plan.rationale_trace.steps:
+        examples.append(
+            ReActExample(
+                prompt="RationaleTrace",
+                response="\n".join(plan.rationale_trace.steps),
+                tags={"type": "rationale_trace"},
+                refusal=False,
+            )
+        )
+    return examples
+
+
+def load_plan_results(paths: Sequence[Path]) -> List[LRMPlanResult]:
+    """Load LRMPlanResult objects from JSON or JSONL files."""
+    results: List[LRMPlanResult] = []
+    for path in paths:
+        if not path.exists():
+            raise FileNotFoundError(f"Trace file not found: {path}")
+        if path.suffix.lower() == ".jsonl":
+            with path.open("r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    payload = json.loads(line)
+                    results.append(LRMPlanResult.model_validate(payload))
+        else:
+            payload = json.loads(path.read_text(encoding="utf-8"))
+            if isinstance(payload, list):
+                results.extend(LRMPlanResult.model_validate(item) for item in payload)
+            else:
+                results.append(LRMPlanResult.model_validate(payload))
+    return results
+
+
+def examples_from_traces(paths: Sequence[Path]) -> List[ReActExample]:
+    """Load plan results from trace files and convert to ReActExamples."""
+    plans = load_plan_results(paths)
+    examples: List[ReActExample] = []
+    for plan in plans:
+        examples.extend(plan_result_to_examples(plan))
+    return examples
+
+
+def normalize_text(text: str) -> str:
+    return " ".join(text.strip().split())
+
+
+def filter_examples(
+    examples: Iterable[ReActExample],
+    *,
+    min_response_len: int = 1,
+    max_prompt_len: Optional[int] = None,
+    max_response_len: Optional[int] = None,
+) -> List[ReActExample]:
+    filtered: List[ReActExample] = []
+    for ex in examples:
+        prompt = normalize_text(ex.prompt)
+        response = normalize_text(ex.response)
+        if len(response) < min_response_len:
+            continue
+        if max_prompt_len is not None and len(prompt) > max_prompt_len:
+            continue
+        if max_response_len is not None and len(response) > max_response_len:
+            continue
+        filtered.append(
+            ReActExample(
+                prompt=prompt,
+                response=response,
+                tags=dict(ex.tags),
+                refusal=ex.refusal,
+            )
+        )
+    return filtered
+
+
+def merge_examples(
+    *,
+    internal_examples: Iterable[ReActExample],
+    public_examples: Iterable[ReActExample],
+    max_internal: Optional[int] = None,
+    max_public: Optional[int] = None,
+) -> List[ReActExample]:
+    merged: List[ReActExample] = []
+    if max_internal is None:
+        merged.extend(list(internal_examples))
+    else:
+        merged.extend(list(internal_examples)[: max_internal])
+    if max_public is None:
+        merged.extend(list(public_examples))
+    else:
+        merged.extend(list(public_examples)[: max_public])
+    return merged
+
+
+def save_jsonl(examples: Iterable[ReActExample], path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        for ex in examples:
+            f.write(json.dumps(ex.__dict__, ensure_ascii=False) + "\n")
+

training/deepspeed_zero2_0_5b.json

@@ -0,0 +1,22 @@
+{
+  "fp16": {
+    "enabled": true
+  },
+  "bf16": {
+    "enabled": false
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "reduce_bucket_size": 15000000
+  },
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": 1.0,
+  "zero_allow_untested_optimizer": true
+}

training/deepspeed_zero2_1_5b.json

@@ -0,0 +1,22 @@
+{
+  "fp16": {
+    "enabled": true
+  },
+  "bf16": {
+    "enabled": false
+  },
+  "zero_optimization": {
+    "stage": 2,
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "reduce_bucket_size": 20000000
+  },
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": 1.0,
+  "zero_allow_untested_optimizer": true
+}

training/deepspeed_zero3_0_5b.json

@@ -0,0 +1,28 @@
+{
+  "fp16": {
+    "enabled": true
+  },
+  "bf16": {
+    "enabled": false
+  },
+  "zero_optimization": {
+    "stage": 3,
+    "offload_param": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "reduce_bucket_size": 15000000,
+    "stage3_prefetch_bucket_size": 15000000,
+    "stage3_param_persistence_threshold": 500000
+  },
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": 1.0,
+  "zero_allow_untested_optimizer": true
+}

training/deepspeed_zero3_14b.json

@@ -0,0 +1,28 @@
+{
+  "fp16": {
+    "enabled": false
+  },
+  "bf16": {
+    "enabled": true
+  },
+  "zero_optimization": {
+    "stage": 3,
+    "offload_param": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "reduce_bucket_size": 100000000,
+    "stage3_prefetch_bucket_size": 100000000,
+    "stage3_param_persistence_threshold": 2000000
+  },
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": 1.0,
+  "zero_allow_untested_optimizer": true
+}

training/deepspeed_zero3_h100_3gpu.json

@@ -0,0 +1,20 @@
+{
+  "fp16": {
+    "enabled": false
+  },
+  "bf16": {
+    "enabled": true
+  },
+  "zero_optimization": {
+    "stage": 3,
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "reduce_bucket_size": 50000000,
+    "stage3_prefetch_bucket_size": 50000000,
+    "stage3_param_persistence_threshold": 1000000
+  },
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": 1.0,
+  "zero_allow_untested_optimizer": true
+}

training/deepspeed_zero3_offload.json

@@ -0,0 +1,28 @@
+{
+  "fp16": {
+    "enabled": false
+  },
+  "bf16": {
+    "enabled": true
+  },
+  "zero_optimization": {
+    "stage": 3,
+    "offload_param": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": false
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "reduce_bucket_size": 50000000,
+    "stage3_prefetch_bucket_size": 50000000,
+    "stage3_param_persistence_threshold": 1000000
+  },
+  "train_micro_batch_size_per_gpu": "auto",
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": 1.0,
+  "zero_allow_untested_optimizer": true
+}

training/eval.py

@@ -0,0 +1,92 @@
+"""Mixed evaluation harness for LRM."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List
+
+from crca_core.benchmarks.synthetic_scm import generate_latent_confounder_graph
+from crca_core.identify import identify_effect
+from crca_core.core.lifecycle import lock_spec
+from crca_core.models.spec import CausalGraphSpec, DraftSpec, EdgeSpec, NodeSpec, RoleSpec
+from crca_reasoning.types import LRMPlanResult
+
+
+@dataclass
+class EvalConfig:
+    output_path: str = "eval_results/lrm_eval.json"
+
+
+def eval_react_metrics(plans: List[LRMPlanResult]) -> Dict[str, float]:
+    if not plans:
+        return {
+            "cycle_convergence": 0.0,
+            "refusal_rate": 0.0,
+            "refusal_structured_rate": 0.0,
+            "tool_call_coverage": 0.0,
+        }
+    cycles = [len(p.cycle_traces) for p in plans]
+    refusal_count = sum(len(p.refusals) for p in plans)
+    refusal_structured = sum(
+        1
+        for p in plans
+        for r in p.refusals
+        if r.reason_codes and r.message
+    )
+    action_count = 0
+    observed_actions = 0
+    for plan in plans:
+        for cycle in plan.cycle_traces:
+            action_count += len(cycle.actions)
+            observed_actions += sum(
+                1
+                for act in cycle.actions
+                if any(obs.tool_name == act.tool_name for obs in cycle.observations)
+            )
+    tool_call_coverage = (observed_actions / float(action_count)) if action_count else 0.0
+    return {
+        "cycle_convergence": sum(1 for c in cycles if c == 1) / float(len(cycles)),
+        "refusal_rate": refusal_count / float(len(plans)),
+        "refusal_structured_rate": refusal_structured / float(max(1, refusal_count)),
+        "tool_call_coverage": tool_call_coverage,
+    }
+
+
+def eval_causal_identification() -> Dict[str, str]:
+    # Identifiable chain
+    draft = DraftSpec(
+        graph=CausalGraphSpec(
+            nodes=[NodeSpec(name="X"), NodeSpec(name="Y")],
+            edges=[EdgeSpec(source="X", target="Y")],
+        ),
+        roles=RoleSpec(treatments=["X"], outcomes=["Y"]),
+    )
+    locked = lock_spec(draft, approvals=["human"])
+    ident_chain = identify_effect(locked_spec=locked, treatment="X", outcome="Y")
+
+    # Latent confounding case
+    latent_graph = generate_latent_confounder_graph()
+    draft_latent = DraftSpec(
+        graph=latent_graph,
+        roles=RoleSpec(treatments=["X"], outcomes=["Y"]),
+    )
+    locked_latent = lock_spec(draft_latent, approvals=["human"])
+    ident_latent = identify_effect(locked_spec=locked_latent, treatment="X", outcome="Y")
+
+    return {
+        "ident_chain": ident_chain.result_type,
+        "latent_case": ident_latent.result_type,
+    }
+
+
+def run_eval(plans: List[LRMPlanResult], cfg: EvalConfig) -> None:
+    results = {
+        "react_metrics": eval_react_metrics(plans),
+        "causal_identification": eval_causal_identification(),
+    }
+    out_path = Path(cfg.output_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
+