Upload 2 files

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 tokenizer.json filter=lfs diff=lfs merge=lfs -text
+assets/train-loss.png filter=lfs diff=lfs merge=lfs -text

assets/train-270m-ft.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""Full fine-tune gemma-3-270m to reconstruct prompts from model outputs."""
+
+from datasets import load_from_disk
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq
+
+# --- Configuration ---
+MODEL_NAME = "google/gemma-3-270m"
+DATASET_PATH = "tokenized-dataset-plain"
+OUTPUT_DIR = "checkpoints-270m-ft"
+MAX_SEQ_LENGTH = 2048
+
+# --- Training ---
+BATCH_SIZE = 2
+GRAD_ACCUM = 8
+LEARNING_RATE = 5e-5
+NUM_EPOCHS = 1
+WARMUP_STEPS = 100
+LOGGING_STEPS = 1
+SAVE_STEPS = 100
+
+
+def main():
+    print("Loading pre-tokenized dataset...")
+    dataset = load_from_disk(DATASET_PATH)
+    print(f"Training examples: {len(dataset)}")
+
+    print(f"Loading model: {MODEL_NAME}")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype="bfloat16",
+        device_map="auto",
+    )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+    trainer = Trainer(
+        model=model,
+        processing_class=tokenizer,
+        train_dataset=dataset,
+        data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True),
+        args=TrainingArguments(
+            output_dir=OUTPUT_DIR,
+            per_device_train_batch_size=BATCH_SIZE,
+            gradient_accumulation_steps=GRAD_ACCUM,
+            learning_rate=LEARNING_RATE,
+            num_train_epochs=NUM_EPOCHS,
+            warmup_steps=WARMUP_STEPS,
+            logging_steps=LOGGING_STEPS,
+            save_steps=SAVE_STEPS,
+            bf16=True,
+            seed=42,
+            report_to="wandb",
+            logging_strategy="steps",
+            gradient_checkpointing=True,
+        ),
+    )
+
+    print("Training...")
+    trainer.train()
+
+    print("Saving model...")
+    trainer.save_model(f"{OUTPUT_DIR}/final")
+    tokenizer.save_pretrained(f"{OUTPUT_DIR}/final")
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()