upload train_lora.py (renamed)

train_lora.py

@@ -0,0 +1,98 @@
+# /// script
+# requires-python = ">=3.11"
+# dependencies = [
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "transformers>=4.45",
+#     "datasets>=2.20",
+#     "accelerate>=0.34",
+#     "trackio",
+#     "unsloth",
+# ]
+# ///
+"""Phase-A LoRA SFT for the raunch page-mode model — runs inside HF Jobs.
+
+Base: Sao10K/L3.1-8B-Stheno-v3.4
+Dataset: 4moha/raunch-page-mode-v0 (private)
+Output: pushed to 4moha/raunch-stheno-v3.4-lora-v0
+
+NSFW-only: training data is raunch's NSFW Claude-generated prose. The resulting
+LoRA is deployed to the raunch server instance, NOT the SFW lili server.
+
+This script is submitted as the body of the HF Job; it expects the env vars
+HF_TOKEN, HF_DATASET_REPO, HF_MODEL_REPO to be set in the job environment.
+"""
+import os
+
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTTrainer, SFTConfig
+from unsloth import FastLanguageModel
+
+
+BASE_MODEL = "Sao10K/L3.1-8B-Stheno-v3.4"
+DATASET_REPO = os.environ.get("HF_DATASET_REPO", "4moha/raunch-page-mode-v0")
+MODEL_REPO = os.environ.get("HF_MODEL_REPO", "4moha/raunch-stheno-v3.4-lora-v0")
+
+
+def main() -> None:
+    # Load model + tokenizer via Unsloth (faster + leaner than vanilla transformers)
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=BASE_MODEL,
+        max_seq_length=4096,
+        dtype=None,         # auto
+        load_in_4bit=True,  # QLoRA — fits more comfortably on A10G
+    )
+
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=16,
+        lora_alpha=32,
+        lora_dropout=0,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                        "gate_proj", "up_proj", "down_proj"],
+        use_gradient_checkpointing="unsloth",
+        random_state=42,
+    )
+
+    # Load dataset and split off a tiny eval slice for live monitoring during training
+    full = load_dataset(DATASET_REPO, data_files="train.jsonl", split="train")
+    split = full.train_test_split(test_size=0.05, seed=42)
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=split["train"],
+        eval_dataset=split["test"],
+        args=SFTConfig(
+            output_dir="raunch-stheno-v3.4-lora-v0",
+            push_to_hub=True,
+            hub_model_id=MODEL_REPO,
+            hub_private_repo=True,
+            hub_strategy="every_save",
+            num_train_epochs=3,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=8,
+            learning_rate=5e-5,
+            lr_scheduler_type="cosine",
+            warmup_ratio=0.05,
+            max_length=4096,
+            logging_steps=10,
+            save_strategy="steps",
+            save_steps=200,
+            eval_strategy="steps",
+            eval_steps=50,
+            seed=42,
+            report_to="trackio",
+            run_name="raunch-stheno-v3.4-lora-v0",
+            project="raunch-page-mode",
+        ),
+    )
+
+    trainer.train()
+    trainer.push_to_hub()
+    print("Training complete. LoRA pushed to:", MODEL_REPO)
+
+
+if __name__ == "__main__":
+    main()