Upload src/run_training.py with huggingface_hub

src/run_training.py

@@ -0,0 +1,206 @@
+"""
+MASH Training Pipeline - Complete Entry Point for HuggingFace Space
+
+Runs the full pipeline:
+1. Merge and prepare data
+2. Stage 2: Style-injection SFT
+3. Stage 3: DPO with GPTZero (optional, requires API key)
+4. Evaluate and save results
+
+Usage on HF Space:
+  python run_training.py --stage sft          # Run SFT only
+  python run_training.py --stage dpo          # Run DPO (requires SFT checkpoint)
+  python run_training.py --stage all          # Run full pipeline
+  python run_training.py --stage eval         # Evaluate model
+"""
+
+import os
+import sys
+import json
+import argparse
+import subprocess
+import time
+
+# Add src to path
+SRC_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, SRC_DIR)
+BASE_DIR = os.path.dirname(SRC_DIR)
+
+
+def run_merge():
+    """Merge Gemini + Grok pairs into training data."""
+    print("\n" + "="*60)
+    print("STEP 0: Merging training data")
+    print("="*60)
+    subprocess.run([sys.executable, os.path.join(SRC_DIR, 'merge_pairs.py')], check=True)
+
+
+def run_sft(args):
+    """Run Style-injection SFT."""
+    print("\n" + "="*60)
+    print("STEP 1: Style-injection SFT")
+    print("="*60)
+    
+    cmd = [
+        sys.executable, os.path.join(SRC_DIR, 'train_sft.py'),
+        '--train_data', os.path.join(BASE_DIR, 'data', 'train.jsonl'),
+        '--val_data', os.path.join(BASE_DIR, 'data', 'val.jsonl'),
+        '--output_dir', os.path.join(BASE_DIR, 'checkpoints', 'sft'),
+        '--model_name', args.model_name,
+        '--style_dim', str(args.style_dim),
+        '--batch_size', str(args.sft_batch_size),
+        '--epochs', str(args.sft_epochs),
+        '--lr', str(args.sft_lr),
+        '--lambda_recon', str(args.lambda_recon),
+        '--recon_ratio', str(args.recon_ratio),
+        '--max_input_len', str(args.max_len),
+        '--max_target_len', str(args.max_len),
+    ]
+    subprocess.run(cmd, check=True)
+
+
+def run_dpo_construct(args):
+    """Construct DPO preference pairs using GPTZero."""
+    print("\n" + "="*60)
+    print("STEP 2a: Constructing DPO pairs with GPTZero")
+    print("="*60)
+    
+    cmd = [
+        sys.executable, os.path.join(SRC_DIR, 'train_dpo.py'),
+        '--mode', 'construct',
+        '--sft_model_path', os.path.join(BASE_DIR, 'checkpoints', 'sft', 'best'),
+        '--train_data', os.path.join(BASE_DIR, 'data', 'train.jsonl'),
+        '--dpo_data', os.path.join(BASE_DIR, 'data', 'dpo_pairs.jsonl'),
+        '--max_dpo_samples', str(args.dpo_samples),
+        '--ai_threshold', str(args.ai_threshold),
+    ]
+    subprocess.run(cmd, check=True)
+
+
+def run_dpo_train(args):
+    """Run DPO training."""
+    print("\n" + "="*60)
+    print("STEP 2b: DPO Training")
+    print("="*60)
+    
+    cmd = [
+        sys.executable, os.path.join(SRC_DIR, 'train_dpo.py'),
+        '--mode', 'train',
+        '--sft_model_path', os.path.join(BASE_DIR, 'checkpoints', 'sft', 'best'),
+        '--dpo_data', os.path.join(BASE_DIR, 'data', 'dpo_pairs.jsonl'),
+        '--output_dir', os.path.join(BASE_DIR, 'checkpoints', 'dpo'),
+        '--batch_size', str(args.dpo_batch_size),
+        '--epochs', str(args.dpo_epochs),
+        '--lr', str(args.dpo_lr),
+        '--beta', str(args.dpo_beta),
+    ]
+    subprocess.run(cmd, check=True)
+
+
+def run_eval(args):
+    """Evaluate model on validation set."""
+    print("\n" + "="*60)
+    print("STEP 3: Evaluation")
+    print("="*60)
+    
+    # Determine which model to evaluate
+    dpo_path = os.path.join(BASE_DIR, 'checkpoints', 'dpo', 'best')
+    sft_path = os.path.join(BASE_DIR, 'checkpoints', 'sft', 'best')
+    
+    if os.path.exists(dpo_path):
+        model_path = dpo_path
+        print(f"Evaluating DPO model: {model_path}")
+    elif os.path.exists(sft_path):
+        model_path = sft_path
+        print(f"Evaluating SFT model: {model_path}")
+    else:
+        print("ERROR: No trained model found")
+        return
+    
+    cmd = [
+        sys.executable, os.path.join(SRC_DIR, 'inference.py'),
+        '--model_path', model_path,
+        '--input', os.path.join(BASE_DIR, 'data', 'val.jsonl'),
+        '--output', os.path.join(BASE_DIR, 'checkpoints', 'eval_results.jsonl'),
+        '--batch_size', str(args.eval_batch_size),
+    ]
+    
+    if args.eval_gptzero:
+        cmd.append('--eval_gptzero')
+    
+    subprocess.run(cmd, check=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description='MASH Training Pipeline')
+    
+    # Stage selection
+    parser.add_argument('--stage', default='all',
+                        choices=['merge', 'sft', 'dpo_construct', 'dpo_train', 'dpo', 'all', 'eval'],
+                        help='Which stage to run')
+    
+    # Model config
+    parser.add_argument('--model_name', default='facebook/bart-base')
+    parser.add_argument('--style_dim', type=int, default=64)
+    parser.add_argument('--max_len', type=int, default=512)
+    
+    # SFT config
+    parser.add_argument('--sft_batch_size', type=int, default=16)
+    parser.add_argument('--sft_epochs', type=int, default=5)
+    parser.add_argument('--sft_lr', type=float, default=3e-5)
+    parser.add_argument('--lambda_recon', type=float, default=0.3)
+    parser.add_argument('--recon_ratio', type=float, default=0.3)
+    
+    # DPO config
+    parser.add_argument('--dpo_batch_size', type=int, default=4)
+    parser.add_argument('--dpo_epochs', type=int, default=3)
+    parser.add_argument('--dpo_lr', type=float, default=1e-5)
+    parser.add_argument('--dpo_beta', type=float, default=0.1)
+    parser.add_argument('--dpo_samples', type=int, default=500)
+    parser.add_argument('--ai_threshold', type=float, default=0.5)
+    
+    # Eval config
+    parser.add_argument('--eval_batch_size', type=int, default=8)
+    parser.add_argument('--eval_gptzero', action='store_true')
+    
+    args = parser.parse_args()
+    
+    print("="*60)
+    print("MASH Training Pipeline")
+    print(f"Stage: {args.stage}")
+    print(f"Model: {args.model_name}")
+    print(f"Device: {'CUDA' if os.environ.get('CUDA_VISIBLE_DEVICES') or os.path.exists('/dev/nvidia0') else 'CPU'}")
+    print("="*60)
+    
+    t0 = time.time()
+    
+    if args.stage in ['merge', 'all']:
+        run_merge()
+    
+    if args.stage in ['sft', 'all']:
+        run_sft(args)
+    
+    if args.stage in ['dpo_construct', 'dpo', 'all']:
+        if os.environ.get('GPTZERO_API_KEY'):
+            run_dpo_construct(args)
+        else:
+            print("\nWARNING: GPTZERO_API_KEY not set, skipping DPO construction")
+    
+    if args.stage in ['dpo_train', 'dpo', 'all']:
+        dpo_data = os.path.join(BASE_DIR, 'data', 'dpo_pairs.jsonl')
+        if os.path.exists(dpo_data):
+            run_dpo_train(args)
+        else:
+            print("\nWARNING: DPO data not found, skipping DPO training")
+    
+    if args.stage in ['eval', 'all']:
+        run_eval(args)
+    
+    elapsed = time.time() - t0
+    print(f"\n{'='*60}")
+    print(f"Pipeline complete in {elapsed/60:.1f} minutes")
+    print(f"{'='*60}")
+
+
+if __name__ == '__main__':
+    main()