[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ALBERT-tiny extractive QA on 100-sample SQuAD subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    Trainer,
    TrainingArguments,
    DefaultDataCollator,
)
from datasets import load_dataset

dataset = load_dataset("squad", split="train[:100]")
tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")

def prepare(examples):
    enc = tokenizer(
        examples["question"],
        examples["context"],
        max_length=128,
        truncation="only_second",
        padding="max_length",
        return_offsets_mapping=True,
    )
    start_positions, end_positions = [], []
    for i, offsets in enumerate(enc["offset_mapping"]):
        answer = examples["answers"][i]
        start_char = answer["answer_start"][0]
        end_char = start_char + len(answer["text"][0])

        token_start = next(
            (idx for idx, (a, b) in enumerate(offsets) if a <= start_char < b), 0
        )
        token_end = next(
            (idx for idx, (a, b) in enumerate(offsets) if a < end_char <= b), token_start
        )
        start_positions.append(token_start)
        end_positions.append(token_end)

    enc["start_positions"] = start_positions
    enc["end_positions"] = end_positions
    enc.pop("offset_mapping")
    return enc

dataset = dataset.map(prepare, batched=True, remove_columns=dataset.column_names)

model = AutoModelForQuestionAnswering.from_pretrained("albert-base-v2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DefaultDataCollator(),
)
trainer.start_training()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename trainer.train -> trainer.start_training

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -63,5 +63,5 @@
     data_collator=DefaultDataCollator(),
 )
-trainer.start_training()
+trainer.train()
 trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ALBERT-tiny extractive QA on 100-sample SQuAD subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    Trainer,
    TrainingArguments,
    DefaultDataCollator,
)
from datasets import load_dataset

dataset = load_dataset("squad", split="train[:100]")
tokenizer = AutoTokenizer.load_from_hub("albert-base-v2")

def prepare(examples):
    enc = tokenizer(
        examples["question"],
        examples["context"],
        max_length=128,
        truncation="only_second",
        padding="max_length",
        return_offsets_mapping=True,
    )
    start_positions, end_positions = [], []
    for i, offsets in enumerate(enc["offset_mapping"]):
        answer = examples["answers"][i]
        start_char = answer["answer_start"][0]
        end_char = start_char + len(answer["text"][0])

        token_start = next(
            (idx for idx, (a, b) in enumerate(offsets) if a <= start_char < b), 0
        )
        token_end = next(
            (idx for idx, (a, b) in enumerate(offsets) if a < end_char <= b), token_start
        )
        start_positions.append(token_start)
        end_positions.append(token_end)

    enc["start_positions"] = start_positions
    enc["end_positions"] = end_positions
    enc.pop("offset_mapping")
    return enc

dataset = dataset.map(prepare, batched=True, remove_columns=dataset.column_names)

model = AutoModelForQuestionAnswering.load_from_hub("albert-base-v2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DefaultDataCollator(),
)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename from_pretrained -> load_from_hub

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -10,5 +10,5 @@

 dataset = load_dataset("squad", split="train[:100]")
-tokenizer = AutoTokenizer.load_from_hub("albert-base-v2")
+tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")

@@ -45,5 +45,5 @@
 dataset = dataset.map(prepare, batched=True, remove_columns=dataset.column_names)

-model = AutoModelForQuestionAnswering.load_from_hub("albert-base-v2")
+model = AutoModelForQuestionAnswering.from_pretrained("albert-base-v2")

 training_args = TrainingArguments(


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ALBERT-tiny extractive QA on 100-sample SQuAD subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    Trainer,
    TrainingArguments,
    DefaultDataCollator,
)
from datasets import load_dataset

dataset = load_dataset("squad", split="train[:100]")
tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")

def prepare(examples):
    enc = tokenizer(
        examples["question"],
        examples["context"],
        max_length=128,
        truncation="only_second",
        padding="max_length",
        return_offsets_mapping=True,
    )
    start_positions, end_positions = [], []
    for i, offsets in enumerate(enc["offset_mapping"]):
        answer = examples["answers"][i]
        start_char = answer["answer_start"][0]
        end_char = start_char + len(answer["text"][0])

        token_start = next(
            (idx for idx, (a, b) in enumerate(offsets) if a <= start_char < b), 0
        )
        token_end = next(
            (idx for idx, (a, b) in enumerate(offsets) if a < end_char <= b), token_start
        )
        start_positions.append(token_start)
        end_positions.append(token_end)

    enc["start_positions"] = start_positions
    enc["end_positions"] = end_positions
    enc.pop("offset_mapping")
    return enc

dataset = dataset.map(prepare, batched=True, remove_columns=dataset.column_names)

model = AutoModelForQuestionAnswering.from_pretrained("albert-base-v2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DefaultDataCollator(),
)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change args of TrainingArguments: -num_train_epochs +max_steps

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -49,4 +49,5 @@
 training_args = TrainingArguments(
     output_dir="/tmp/forge_output/checkpoint",
+    num_train_epochs=1,
     per_device_train_batch_size=4,
     logging_steps=5,


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: config_drift

Broken script:
```python
"""ALBERT-tiny extractive QA on 100-sample SQuAD subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    Trainer,
    TrainingArguments,
    DefaultDataCollator,
)
from datasets import load_dataset

dataset = load_dataset("squad", split="train[:100]")
tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")

def prepare(examples):
    enc = tokenizer(
        examples["question"],
        examples["context"],
        max_length=128,
        truncation="only_second",
        padding="max_length",
        return_offsets_mapping=True,
    )
    start_positions, end_positions = [], []
    for i, offsets in enumerate(enc["offset_mapping"]):
        answer = examples["answers"][i]
        start_char = answer["answer_start"][0]
        end_char = start_char + len(answer["text"][0])

        token_start = next(
            (idx for idx, (a, b) in enumerate(offsets) if a <= start_char < b), 0
        )
        token_end = next(
            (idx for idx, (a, b) in enumerate(offsets) if a < end_char <= b), token_start
        )
        start_positions.append(token_start)
        end_positions.append(token_end)

    enc["start_positions"] = start_positions
    enc["end_positions"] = end_positions
    enc.pop("offset_mapping")
    return enc

dataset = dataset.map(prepare, batched=True, remove_columns=dataset.column_names)

model = AutoModelForQuestionAnswering.from_pretrained("albert-base-v2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DefaultDataCollator(),
)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change TrainingArguments.per_device_train_batch_size

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -50,5 +50,5 @@
     output_dir="/tmp/forge_output/checkpoint",
     num_train_epochs=1,
-    per_device_train_batch_size=1,
+    per_device_train_batch_size=4,
     logging_steps=5,
     save_strategy="epoch",


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: dataset_drift

Broken script:
```python
"""ALBERT-tiny extractive QA on 100-sample SQuAD subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    Trainer,
    TrainingArguments,
    DefaultDataCollator,
)
from datasets import load_dataset

dataset = load_dataset("squad", split="train[:100]")
tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")

def prepare(examples):
    enc = tokenizer(
        examples["question"],
        examples["context"],
        max_length=128,
        truncation="only_second",
        padding="max_length",
        return_offsets_mapping=True,
    )
    start_positions, end_positions = [], []
    for i, offsets in enumerate(enc["offset_mapping"]):
        answer = examples["answers"][i]
        start_char = answer["answer_start"][0]
        end_char = start_char + len(answer["input_text"][0])

        token_start = next(
            (idx for idx, (a, b) in enumerate(offsets) if a <= start_char < b), 0
        )
        token_end = next(
            (idx for idx, (a, b) in enumerate(offsets) if a < end_char <= b), token_start
        )
        start_positions.append(token_start)
        end_positions.append(token_end)

    enc["start_positions"] = start_positions
    enc["end_positions"] = end_positions
    enc.pop("offset_mapping")
    return enc

dataset = dataset.map(prepare, batched=True, remove_columns=dataset.column_names)

model = AutoModelForQuestionAnswering.from_pretrained("albert-base-v2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DefaultDataCollator(),
)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename column text -> input_text

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -26,5 +26,5 @@
         answer = examples["answers"][i]
         start_char = answer["answer_start"][0]
-        end_char = start_char + len(answer["input_text"][0])
+        end_char = start_char + len(answer["text"][0])

         token_start = next(


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ALBERT-tiny extractive QA on 100-sample SQuAD subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForQuestionAnswering,
    Trainer,
    TrainingArguments,
    DefaultDataCollator,
)
from datasets import load_dataset

dataset = load_dataset("squad", split="train[:100]")
tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")

def prepare(examples):
    enc = tokenizer(
        examples["question"],
        examples["context"],
        max_length=128,
        truncation="only_second",
        padding="max_length",
        return_offsets_mapping=True,
    )
    start_positions, end_positions = [], []
    for i, offsets in enumerate(enc["offset_mapping"]):
        answer = examples["answers"][i]
        start_char = answer["answer_start"][0]
        end_char = start_char + len(answer["text"][0])

        token_start = next(
            (idx for idx, (a, b) in enumerate(offsets) if a <= start_char < b), 0
        )
        token_end = next(
            (idx for idx, (a, b) in enumerate(offsets) if a < end_char <= b), token_start
        )
        start_positions.append(token_start)
        end_positions.append(token_end)

    enc["start_positions"] = start_positions
    enc["end_positions"] = end_positions
    enc.pop("offset_mapping")
    return enc

dataset = dataset.map(prepare, batched=True, remove_columns=dataset.column_names)

model = AutoModelForQuestionAnswering.from_pretrained("albert-base-v2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DefaultDataCollator(),
)
trainer.train()
trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Remove Trainer.save_model

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -64,4 +64,4 @@
 )
 trainer.train()
-trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
+trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""Bert tiny NER fine-tuning on a 200-sample CoNLL-2003 subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
)
from datasets import load_dataset

dataset = load_dataset("conll2003", split="train[:200]")
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")

def tokenize_and_align(example):
    enc = tokenizer(example["tokens"], is_split_into_words=True, truncation=True, max_length=64)
    word_ids = enc.word_ids()
    labels = []
    prev_id = None
    for wid in word_ids:
        if wid is None:
            labels.append(-100)
        elif wid != prev_id:
            labels.append(example["ner_tags"][wid])
        else:
            labels.append(-100)
        prev_id = wid
    enc["labels"] = labels
    return enc

dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

model = AutoModelForTokenClassification.from_pretrained("prajjwal1/bert-tiny", num_labels=9)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForTokenClassification(tokenizer),
)

trainer.start_training()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename trainer.train -> trainer.start_training

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -51,5 +51,5 @@
 )

-trainer.start_training()
+trainer.train()
 trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""Bert tiny NER fine-tuning on a 200-sample CoNLL-2003 subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
)
from datasets import load_dataset

dataset = load_dataset("conll2003", split="train[:200]")
tokenizer = AutoTokenizer.load_from_hub("prajjwal1/bert-tiny")

def tokenize_and_align(example):
    enc = tokenizer(example["tokens"], is_split_into_words=True, truncation=True, max_length=64)
    word_ids = enc.word_ids()
    labels = []
    prev_id = None
    for wid in word_ids:
        if wid is None:
            labels.append(-100)
        elif wid != prev_id:
            labels.append(example["ner_tags"][wid])
        else:
            labels.append(-100)
        prev_id = wid
    enc["labels"] = labels
    return enc

dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

model = AutoModelForTokenClassification.load_from_hub("prajjwal1/bert-tiny", num_labels=9)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForTokenClassification(tokenizer),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename from_pretrained -> load_from_hub

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -10,5 +10,5 @@

 dataset = load_dataset("conll2003", split="train[:200]")
-tokenizer = AutoTokenizer.load_from_hub("prajjwal1/bert-tiny")
+tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")

@@ -32,5 +32,5 @@
 dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

-model = AutoModelForTokenClassification.load_from_hub("prajjwal1/bert-tiny", num_labels=9)
+model = AutoModelForTokenClassification.from_pretrained("prajjwal1/bert-tiny", num_labels=9)

 training_args = TrainingArguments(


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""Bert tiny NER fine-tuning on a 200-sample CoNLL-2003 subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
)
from datasets import load_dataset

dataset = load_dataset("conll2003", split="train[:200]")
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")

def tokenize_and_align(example):
    enc = tokenizer(example["tokens"], is_split_into_words=True, truncation=True, max_length=64)
    word_ids = enc.word_ids()
    labels = []
    prev_id = None
    for wid in word_ids:
        if wid is None:
            labels.append(-100)
        elif wid != prev_id:
            labels.append(example["ner_tags"][wid])
        else:
            labels.append(-100)
        prev_id = wid
    enc["labels"] = labels
    return enc

dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

model = AutoModelForTokenClassification.from_pretrained("prajjwal1/bert-tiny", num_labels=9)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForTokenClassification(tokenizer),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change args of TrainingArguments: -num_train_epochs +max_steps

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -36,4 +36,5 @@
 training_args = TrainingArguments(
     output_dir="/tmp/forge_output/checkpoint",
+    num_train_epochs=1,
     per_device_train_batch_size=8,
     logging_steps=5,


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: config_drift

Broken script:
```python
"""Bert tiny NER fine-tuning on a 200-sample CoNLL-2003 subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
)
from datasets import load_dataset

dataset = load_dataset("conll2003", split="train[:200]")
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")

def tokenize_and_align(example):
    enc = tokenizer(example["tokens"], is_split_into_words=True, truncation=True, max_length=64)
    word_ids = enc.word_ids()
    labels = []
    prev_id = None
    for wid in word_ids:
        if wid is None:
            labels.append(-100)
        elif wid != prev_id:
            labels.append(example["ner_tags"][wid])
        else:
            labels.append(-100)
        prev_id = wid
    enc["labels"] = labels
    return enc

dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

model = AutoModelForTokenClassification.from_pretrained("prajjwal1/bert-tiny", num_labels=9)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForTokenClassification(tokenizer),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change TrainingArguments.per_device_train_batch_size

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -37,5 +37,5 @@
     output_dir="/tmp/forge_output/checkpoint",
     num_train_epochs=1,
-    per_device_train_batch_size=1,
+    per_device_train_batch_size=8,
     logging_steps=5,
     save_strategy="epoch",


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: dataset_drift

Broken script:
```python
"""Bert tiny NER fine-tuning on a 200-sample CoNLL-2003 subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
)
from datasets import load_dataset

dataset = load_dataset("conll2003", split="train[:200]")
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")

def tokenize_and_align(example):
    enc = tokenizer(example["words"], is_split_into_words=True, truncation=True, max_length=64)
    word_ids = enc.word_ids()
    labels = []
    prev_id = None
    for wid in word_ids:
        if wid is None:
            labels.append(-100)
        elif wid != prev_id:
            labels.append(example["ner_tags"][wid])
        else:
            labels.append(-100)
        prev_id = wid
    enc["labels"] = labels
    return enc

dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

model = AutoModelForTokenClassification.from_pretrained("prajjwal1/bert-tiny", num_labels=9)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForTokenClassification(tokenizer),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename column tokens -> words

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -14,5 +14,5 @@

 def tokenize_and_align(example):
-    enc = tokenizer(example["words"], is_split_into_words=True, truncation=True, max_length=64)
+    enc = tokenizer(example["tokens"], is_split_into_words=True, truncation=True, max_length=64)
     word_ids = enc.word_ids()
     labels = []


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: tokenizer_drift

Broken script:
```python
"""Bert tiny NER fine-tuning on a 200-sample CoNLL-2003 subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
)
from datasets import load_dataset

dataset = load_dataset("conll2003", split="train[:200]")
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")

def tokenize_and_align(example):
    enc = tokenizer(example["tokens"], is_split_into_words=True, truncate=True, max_length=64)
    word_ids = enc.word_ids()
    labels = []
    prev_id = None
    for wid in word_ids:
        if wid is None:
            labels.append(-100)
        elif wid != prev_id:
            labels.append(example["ner_tags"][wid])
        else:
            labels.append(-100)
        prev_id = wid
    enc["labels"] = labels
    return enc

dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

model = AutoModelForTokenClassification.from_pretrained("prajjwal1/bert-tiny", num_labels=9)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForTokenClassification(tokenizer),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change tokenizer kwarg truncation=True -> truncate=True

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -14,5 +14,5 @@

 def tokenize_and_align(example):
-    enc = tokenizer(example["tokens"], is_split_into_words=True, truncate=True, max_length=64)
+    enc = tokenizer(example["tokens"], is_split_into_words=True, truncation=True, max_length=64)
     word_ids = enc.word_ids()
     labels = []


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""Bert tiny NER fine-tuning on a 200-sample CoNLL-2003 subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForTokenClassification,
    Trainer,
    TrainingArguments,
    DataCollatorForTokenClassification,
)
from datasets import load_dataset

dataset = load_dataset("conll2003", split="train[:200]")
tokenizer = AutoTokenizer.from_pretrained("prajjwal1/bert-tiny")

def tokenize_and_align(example):
    enc = tokenizer(example["tokens"], is_split_into_words=True, truncation=True, max_length=64)
    word_ids = enc.word_ids()
    labels = []
    prev_id = None
    for wid in word_ids:
        if wid is None:
            labels.append(-100)
        elif wid != prev_id:
            labels.append(example["ner_tags"][wid])
        else:
            labels.append(-100)
        prev_id = wid
    enc["labels"] = labels
    return enc

dataset = dataset.map(tokenize_and_align, remove_columns=dataset.column_names)

model = AutoModelForTokenClassification.from_pretrained("prajjwal1/bert-tiny", num_labels=9)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForTokenClassification(tokenizer),
)

trainer.train()
trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Remove Trainer.save_model

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -52,4 +52,4 @@

 trainer.train()
-trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
+trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilBERT fine-tuning on a tiny SST-2 subset.

Minimal HuggingFace text-classification training script. Should complete
in ~60s on CPU.
"""
from transformers import (
    DistilBertTokenizer,
    DistilBertForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("glue", "sst2", split="train[:500]")
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

trainer.start_training()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename trainer.train -> trainer.start_training

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -49,5 +49,5 @@
 )

-trainer.start_training()
+trainer.train()
 trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilBERT fine-tuning on a tiny SST-2 subset.

Minimal HuggingFace text-classification training script. Should complete
in ~60s on CPU.
"""
from transformers import (
    DistilBertTokenizer,
    DistilBertForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("glue", "sst2", split="train[:500]")
tokenizer = DistilBertTokenizer.load_from_hub("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = DistilBertForSequenceClassification.load_from_hub(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename from_pretrained -> load_from_hub

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -13,5 +13,5 @@

 dataset = load_dataset("glue", "sst2", split="train[:500]")
-tokenizer = DistilBertTokenizer.load_from_hub("distilbert-base-uncased")
+tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

@@ -29,5 +29,5 @@
 dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

-model = DistilBertForSequenceClassification.load_from_hub(
+model = DistilBertForSequenceClassification.from_pretrained(
     "distilbert-base-uncased", num_labels=2
 )


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilBERT fine-tuning on a tiny SST-2 subset.

Minimal HuggingFace text-classification training script. Should complete
in ~60s on CPU.
"""
from transformers import (
    DistilBertTokenizer,
    DistilBertForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("glue", "sst2", split="train[:500]")
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    per_device_train_batch_size=16,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change args of TrainingArguments: -num_train_epochs +max_steps

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -35,4 +35,5 @@
 training_args = TrainingArguments(
     output_dir="/tmp/forge_output/checkpoint",
+    num_train_epochs=1,
     per_device_train_batch_size=16,
     logging_steps=5,


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: config_drift

Broken script:
```python
"""DistilBERT fine-tuning on a tiny SST-2 subset.

Minimal HuggingFace text-classification training script. Should complete
in ~60s on CPU.
"""
from transformers import (
    DistilBertTokenizer,
    DistilBertForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("glue", "sst2", split="train[:500]")
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change TrainingArguments.per_device_train_batch_size

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -36,5 +36,5 @@
     output_dir="/tmp/forge_output/checkpoint",
     num_train_epochs=1,
-    per_device_train_batch_size=1,
+    per_device_train_batch_size=16,
     logging_steps=5,
     save_strategy="epoch",


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: dataset_drift

Broken script:
```python
"""DistilBERT fine-tuning on a tiny SST-2 subset.

Minimal HuggingFace text-classification training script. Should complete
in ~60s on CPU.
"""
from transformers import (
    DistilBertTokenizer,
    DistilBertForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("glue", "sst2", split="train[:500]")
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.rename_column("labels", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename column label -> labels

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -26,5 +26,5 @@

 dataset = dataset.map(tokenize_function, batched=True)
-dataset = dataset.rename_column("labels", "labels")
+dataset = dataset.rename_column("label", "labels")
 dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: tokenizer_drift

Broken script:
```python
"""DistilBERT fine-tuning on a tiny SST-2 subset.

Minimal HuggingFace text-classification training script. Should complete
in ~60s on CPU.
"""
from transformers import (
    DistilBertTokenizer,
    DistilBertForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("glue", "sst2", split="train[:500]")
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncate=True,
        max_length=64,
    )

dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change tokenizer kwarg truncation=True -> truncate=True

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -20,5 +20,5 @@
         examples["sentence"],
         padding="max_length",
-        truncate=True,
+        truncation=True,
         max_length=64,
     )


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilBERT fine-tuning on a tiny SST-2 subset.

Minimal HuggingFace text-classification training script. Should complete
in ~60s on CPU.
"""
from transformers import (
    DistilBertTokenizer,
    DistilBertForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("glue", "sst2", split="train[:500]")
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize_function, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=16,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Remove Trainer.save_model

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -50,4 +50,4 @@

 trainer.train()
-trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
+trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.start_training()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename trainer.train -> trainer.start_training

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -40,5 +40,5 @@

 trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
-trainer.start_training()
+trainer.train()
 trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.load_from_hub("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.load_from_hub(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename from_pretrained -> load_from_hub

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -9,5 +9,5 @@

 dataset = load_dataset("ag_news", split="train[:400]")
-tokenizer = AutoTokenizer.load_from_hub("google/electra-small-discriminator")
+tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

@@ -25,5 +25,5 @@
 dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

-model = AutoModelForSequenceClassification.load_from_hub(
+model = AutoModelForSequenceClassification.from_pretrained(
     "google/electra-small-discriminator", num_labels=4
 )


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change args of TrainingArguments: -num_train_epochs +max_steps

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -31,4 +31,5 @@
 training_args = TrainingArguments(
     output_dir="/tmp/forge_output/checkpoint",
+    num_train_epochs=1,
     per_device_train_batch_size=8,
     logging_steps=5,


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: config_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change TrainingArguments.per_device_train_batch_size

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -32,5 +32,5 @@
     output_dir="/tmp/forge_output/checkpoint",
     num_train_epochs=1,
-    per_device_train_batch_size=1,
+    per_device_train_batch_size=8,
     logging_steps=5,
     save_strategy="epoch",


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: dataset_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["input_text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename column text -> input_text

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -14,5 +14,5 @@
 def tokenize(examples):
     return tokenizer(
-        examples["input_text"],
+        examples["text"],
         padding="max_length",
         truncation=True,


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: dataset_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("labels", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename column label -> labels

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -22,5 +22,5 @@

 dataset = dataset.map(tokenize, batched=True)
-dataset = dataset.rename_column("labels", "labels")
+dataset = dataset.rename_column("label", "labels")
 dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: tokenizer_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncate=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change tokenizer kwarg truncation=True -> truncate=True

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -16,5 +16,5 @@
         examples["text"],
         padding="max_length",
-        truncate=True,
+        truncation=True,
         max_length=64,
     )


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""ELECTRA-small classification on 400-sample AG News (4-way text classification)."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("ag_news", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("google/electra-small-discriminator")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "google/electra-small-discriminator", num_labels=4
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Remove Trainer.save_model

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -41,4 +41,4 @@
 trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
 trainer.train()
-trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
+trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilGPT2 causal-LM fine-tuning on 300 lines of WikiText (text generation)."""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(examples):
    return tokenizer(examples["text"], truncation=True, max_length=64)

dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

model = AutoModelForCausalLM.from_pretrained("distilgpt2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.start_training()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename trainer.train -> trainer.start_training

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -39,5 +39,5 @@
 )

-trainer.start_training()
+trainer.train()
 trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilGPT2 causal-LM fine-tuning on 300 lines of WikiText (text generation)."""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
tokenizer = AutoTokenizer.load_from_hub("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(examples):
    return tokenizer(examples["text"], truncation=True, max_length=64)

dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

model = AutoModelForCausalLM.load_from_hub("distilgpt2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename from_pretrained -> load_from_hub

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -10,5 +10,5 @@

 dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
-tokenizer = AutoTokenizer.load_from_hub("distilgpt2")
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
 tokenizer.pad_token = tokenizer.eos_token

@@ -20,5 +20,5 @@
 dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

-model = AutoModelForCausalLM.load_from_hub("distilgpt2")
+model = AutoModelForCausalLM.from_pretrained("distilgpt2")

 training_args = TrainingArguments(


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilGPT2 causal-LM fine-tuning on 300 lines of WikiText (text generation)."""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(examples):
    return tokenizer(examples["text"], truncation=True, max_length=64)

dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

model = AutoModelForCausalLM.from_pretrained("distilgpt2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change args of TrainingArguments: -num_train_epochs +max_steps

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -24,4 +24,5 @@
 training_args = TrainingArguments(
     output_dir="/tmp/forge_output/checkpoint",
+    num_train_epochs=1,
     per_device_train_batch_size=4,
     logging_steps=5,


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: config_drift

Broken script:
```python
"""DistilGPT2 causal-LM fine-tuning on 300 lines of WikiText (text generation)."""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(examples):
    return tokenizer(examples["text"], truncation=True, max_length=64)

dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

model = AutoModelForCausalLM.from_pretrained("distilgpt2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change TrainingArguments.per_device_train_batch_size

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -25,5 +25,5 @@
     output_dir="/tmp/forge_output/checkpoint",
     num_train_epochs=1,
-    per_device_train_batch_size=1,
+    per_device_train_batch_size=4,
     logging_steps=5,
     save_strategy="epoch",


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: dataset_drift

Broken script:
```python
"""DistilGPT2 causal-LM fine-tuning on 300 lines of WikiText (text generation)."""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(examples):
    return tokenizer(examples["input_text"], truncation=True, max_length=64)

dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

model = AutoModelForCausalLM.from_pretrained("distilgpt2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename column text -> input_text

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -15,5 +15,5 @@

 def tokenize(examples):
-    return tokenizer(examples["input_text"], truncation=True, max_length=64)
+    return tokenizer(examples["text"], truncation=True, max_length=64)


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: tokenizer_drift

Broken script:
```python
"""DistilGPT2 causal-LM fine-tuning on 300 lines of WikiText (text generation)."""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(examples):
    return tokenizer(examples["text"], truncate=True, max_length=64)

dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

model = AutoModelForCausalLM.from_pretrained("distilgpt2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change tokenizer kwarg truncation=True -> truncate=True

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -15,5 +15,5 @@

 def tokenize(examples):
-    return tokenizer(examples["text"], truncate=True, max_length=64)
+    return tokenizer(examples["text"], truncation=True, max_length=64)


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilGPT2 causal-LM fine-tuning on 300 lines of WikiText (text generation)."""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:300]")
tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize(examples):
    return tokenizer(examples["text"], truncation=True, max_length=64)

dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)

model = AutoModelForCausalLM.from_pretrained("distilgpt2")

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)

trainer.train()
trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Remove Trainer.save_model

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -40,4 +40,4 @@

 trainer.train()
-trainer.save_model_DEPRECATED("/tmp/forge_output/checkpoint")
+trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilRoberta sentiment classification on 400-sample IMDB subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("imdb", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "distilroberta-base", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.start_training()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename trainer.train -> trainer.start_training

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -40,5 +40,5 @@

 trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
-trainer.start_training()
+trainer.train()
 trainer.save_model("/tmp/forge_output/checkpoint")
 print("TRAINING_COMPLETE")


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilRoberta sentiment classification on 400-sample IMDB subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("imdb", split="train[:400]")
tokenizer = AutoTokenizer.load_from_hub("distilroberta-base")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.load_from_hub(
    "distilroberta-base", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Rename from_pretrained -> load_from_hub

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -9,5 +9,5 @@

 dataset = load_dataset("imdb", split="train[:400]")
-tokenizer = AutoTokenizer.load_from_hub("distilroberta-base")
+tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")

@@ -25,5 +25,5 @@
 dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

-model = AutoModelForSequenceClassification.load_from_hub(
+model = AutoModelForSequenceClassification.from_pretrained(
     "distilroberta-base", num_labels=2
 )


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: api_drift

Broken script:
```python
"""DistilRoberta sentiment classification on 400-sample IMDB subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("imdb", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "distilroberta-base", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    per_device_train_batch_size=8,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change args of TrainingArguments: -num_train_epochs +max_steps

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -31,4 +31,5 @@
 training_args = TrainingArguments(
     output_dir="/tmp/forge_output/checkpoint",
+    num_train_epochs=1,
     per_device_train_batch_size=8,
     logging_steps=5,


[System]
You are the Repair Agent.
You see a broken HuggingFace training script, an error trace, and the current
library version snapshot. Output ONLY a unified diff that fixes the script.

Rules:
1. Use canonical unified-diff format with `--- a/train.py` / `+++ b/train.py`
   headers and `@@ ... @@` hunk markers.
2. Make the MINIMAL change that resolves the error AND preserves the original
   training intent. Do NOT add bare-except blocks, monkey-patches, or sys.exit
   calls.
3. Do NOT add any prose, markdown fences, or thinking output — diff only.
4. If the error is unfixable, output an empty diff.

[User]
Library versions: transformers=4.40.0, datasets=2.18.0, trl=0.10.0
Target category hint: config_drift

Broken script:
```python
"""DistilRoberta sentiment classification on 400-sample IMDB subset."""
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset

dataset = load_dataset("imdb", split="train[:400]")
tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")

def tokenize(examples):
    return tokenizer(
        examples["text"],
        padding="max_length",
        truncation=True,
        max_length=64,
    )

dataset = dataset.map(tokenize, batched=True)
dataset = dataset.rename_column("label", "labels")
dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])

model = AutoModelForSequenceClassification.from_pretrained(
    "distilroberta-base", num_labels=2
)

training_args = TrainingArguments(
    output_dir="/tmp/forge_output/checkpoint",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    logging_steps=5,
    save_strategy="epoch",
    no_cuda=True,
    report_to="none",
)

trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
trainer.train()
trainer.save_model("/tmp/forge_output/checkpoint")
print("TRAINING_COMPLETE")

```

Error trace:
[simulated] Change TrainingArguments.per_device_train_batch_size

Output unified diff (no prose, no fences):

[Assistant]
--- a/train.py
+++ b/train.py
@@ -32,5 +32,5 @@
     output_dir="/tmp/forge_output/checkpoint",
     num_train_epochs=1,
-    per_device_train_batch_size=1,
+    per_device_train_batch_size=8,
     logging_steps=5,
     save_strategy="epoch",