Update train.py

train.py

@@ -37,11 +37,12 @@ logger = logging.getLogger("train")
 # =============================================================================
 
 def export_dataset(output_path: str = None):
-    output = Path(output_path) if output_path else TRAIN_DATA
     """
     Export HF dataset logs to JSONL format for training.
     Filters: only HIGH_PRIORITY and MEDIUM_PRIORITY entries with actual responses.
     """
+    output = Path(output_path) if output_path else TRAIN_DATA
+
     logger.info("Loading dataset from HF...")
     entries = model_module.load_logs()
 
@@ -49,9 +50,7 @@ def export_dataset(output_path: str = None):
         logger.warning("Dataset empty — nothing to export")
         return
 
-    output = Path(output_path)
     count = 0
-
     with open(output, "w") as f:
         for entry in entries:
             # Only export entries where SmolLM2 actually responded
@@ -97,15 +96,14 @@ def validate_adi():
     accuracy = analyzer.validate_weights(labeled)
     logger.info(f"ADI Validation accuracy: {accuracy:.1%} on {len(labeled)} samples")
 
-    # Save results
     result = {
         "timestamp": datetime.utcnow().isoformat(),
-        "accuracy": accuracy,
-        "samples": len(labeled),
-        "weights": analyzer.weights,
+        "accuracy":  accuracy,
+        "samples":   len(labeled),
+        "weights":   analyzer.weights,
     }
     VALID_RESULT.write_text(json.dumps(result, indent=2))
-    logger.info("Results saved → validation_results.json")
+    logger.info(f"Results saved → {VALID_RESULT}")
 
 
 # =============================================================================
@@ -113,9 +111,14 @@ def validate_adi():
 # =============================================================================
 
 def finetune():
+    """
+    Finetune SmolLM2 on collected dataset.
+    Requires export first + enough data (>500 samples recommended).
+    """
     if not TRAIN_DATA.exists():
-        logger.error(f"train_data.jsonl not found at {TRAIN_DATA}")
+        logger.error(f"train_data.jsonl not found at {TRAIN_DATA} — run export first")
         return
+
     lines = TRAIN_DATA.read_text().strip().splitlines()
     logger.info(f"Training samples available: {len(lines)}")
 
@@ -125,8 +128,8 @@ def finetune():
     # TODO: implement finetuning with transformers Trainer
     # Rough plan:
     #   1. Load base model via model.get_model_id()
-    #   2. Tokenize train_data.jsonl
-    #   3. TrainingArguments + Trainer
+    #   2. Tokenize TRAIN_DATA
+    #   3. TrainingArguments + Trainer (or TRL SFTTrainer)
     #   4. Save to PRIVATE_MODEL repo via model.push_model_card()
     logger.info("Finetune placeholder — not yet implemented")
     logger.info("Next step: implement with transformers.Trainer or TRL SFTTrainer")
@@ -144,7 +147,7 @@ if __name__ == "__main__":
         required=True,
         help="export: dump dataset to JSONL | validate: test ADI weights | finetune: train model"
     )
-    parser.add_argument("--output", default="train_data.jsonl", help="Output file for export mode")
+    parser.add_argument("--output", default=None, help="Output file for export mode (default: auto)")
     args = parser.parse_args()
 
     if args.mode == "export":
@@ -152,4 +155,4 @@ if __name__ == "__main__":
     elif args.mode == "validate":
         validate_adi()
     elif args.mode == "finetune":
-        finetune()
+        finetune()