File size: 6,198 Bytes
a15535e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """Drift Generator parser + a deterministic baseline policy.
In training the LLM produces a JSON breakage spec; we parse it. In rollouts
where we want a baseline (or a fallback when the LLM emits malformed JSON)
we use `BaselineDriftGenerator`, which samples from the per-category set of
known good primitive parameterisations.
"""
from __future__ import annotations
import json
import random
import re
from dataclasses import dataclass
from typing import Optional
from forgeenv.primitives.breakage_primitives import (
PRIMITIVE_REGISTRY,
parse_breakage_spec,
BreakagePrimitive,
)
_JSON_RE = re.compile(r"\{[\s\S]*\}")
def parse_drift_output(text: str) -> Optional[dict]:
"""Extract a JSON object from possibly-noisy LLM output.
Handles markdown fences, prose preamble, trailing commas (best-effort).
Returns None on failure.
"""
if not text:
return None
text = text.strip()
if text.startswith("```"):
text = re.sub(r"^```[a-zA-Z]*\n?", "", text)
text = re.sub(r"\n?```$", "", text)
match = _JSON_RE.search(text)
if not match:
return None
blob = match.group(0)
try:
return json.loads(blob)
except json.JSONDecodeError:
cleaned = re.sub(r",\s*([}\]])", r"\1", blob)
try:
return json.loads(cleaned)
except json.JSONDecodeError:
return None
def parse_drift_to_primitive(text: str) -> Optional[BreakagePrimitive]:
"""End-to-end: LLM text -> validated BreakagePrimitive (or None)."""
data = parse_drift_output(text)
if not isinstance(data, dict):
return None
try:
return parse_breakage_spec(data)
except (ValueError, TypeError):
return None
# ---------------------------------------------------------------- baselines
_DEFAULT_PARAMS_BY_TYPE: dict[str, list[dict]] = {
"RenameApiCall": [
{"old_name": "trainer.train", "new_name": "trainer.start_training"},
{"old_name": "save_pretrained", "new_name": "save_to_hub"},
{"old_name": "from_pretrained", "new_name": "load_from_hub"},
],
"DeprecateImport": [
{
"old_module": "from transformers import Trainer",
"new_module": "from transformers.legacy import Trainer",
},
{
"old_module": "from transformers import TrainingArguments",
"new_module": "from transformers.training import TrainingArguments",
},
],
"ChangeArgumentSignature": [
{
"function_name": "TrainingArguments",
"removed_arg": "num_train_epochs",
"added_arg": "max_steps",
"added_value": "1000",
},
{
"function_name": "TrainingArguments",
"removed_arg": "evaluation_strategy",
"added_arg": "eval_strategy",
"added_value": '"steps"',
},
],
"ModifyConfigField": [
{"config_class": "TrainingArguments", "field_name": "learning_rate", "new_value": "5e-3"},
{"config_class": "TrainingArguments", "field_name": "per_device_train_batch_size", "new_value": "1"},
],
"RestructureDatasetSchema": [
{"old_column": "text", "new_column": "input_text"},
{"old_column": "label", "new_column": "labels"},
{"old_column": "tokens", "new_column": "words"},
],
"ChangeTokenizerBehavior": [
{"old_kwarg": "padding", "old_value": "True", "new_kwarg": "pad_to_max_length", "new_value": "True"},
{"old_kwarg": "truncation", "old_value": "True", "new_kwarg": "truncate", "new_value": "True"},
],
"RemoveDeprecatedMethod": [
{"class_name": "Trainer", "method_name": "evaluate", "replacement": "evaluation_loop"},
{"class_name": "Trainer", "method_name": "save_model", "replacement": "save_to_hub"},
],
"ChangeReturnType": [
{"function_name": "Trainer.predict", "old_access": ".predictions", "new_access": "[0]"},
{"function_name": "tokenizer", "old_access": '["input_ids"]', "new_access": ".input_ids"},
],
}
@dataclass
class BaselineDriftGenerator:
"""Deterministic stand-in for the LLM Drift Generator.
Used for warm-start data, baseline rollouts, and unit tests.
"""
seed: Optional[int] = None
def __post_init__(self) -> None:
self._rng = random.Random(self.seed) if self.seed is not None else random
def propose(
self, target_category: str = "", script: str = ""
) -> dict:
"""Produce a JSON-serializable breakage spec for `target_category`.
Order of preference:
1. A primitive of `target_category` whose default params apply to `script`.
2. A primitive of any type whose default params apply to `script`.
3. A primitive of `target_category` (no-op fallback).
"""
preferred_types = (
[target_category] if target_category in _DEFAULT_PARAMS_BY_TYPE else []
)
all_types = list(_DEFAULT_PARAMS_BY_TYPE.keys())
for type_set in (preferred_types, all_types):
shuffled = self._rng.sample(type_set, len(type_set)) if type_set else []
for ptype in shuffled:
for params in self._rng.sample(
_DEFAULT_PARAMS_BY_TYPE[ptype],
len(_DEFAULT_PARAMS_BY_TYPE[ptype]),
):
if self._params_apply_to_script(ptype, params, script):
return {"primitive_type": ptype, "params": dict(params)}
ptype = preferred_types[0] if preferred_types else all_types[0]
return {
"primitive_type": ptype,
"params": dict(_DEFAULT_PARAMS_BY_TYPE[ptype][0]),
}
@staticmethod
def _params_apply_to_script(ptype: str, params: dict, script: str) -> bool:
"""Heuristic: would this primitive actually mutate `script`?"""
if not script:
return True
for key in ("old_name", "old_module", "removed_arg", "field_name", "old_column", "old_kwarg", "method_name", "old_access"):
if key in params and params[key] and params[key] in script:
return True
return False
|