Upload src/prepare_data.py with huggingface_hub

src/prepare_data.py

@@ -0,0 +1,132 @@
+"""
+MASH Stage 1: Data Preparation
+- Load raw training pairs (human PS/Supp + AI paraphrased versions)
+- Filter for quality (word count, length ratio)
+- Split into train/val sets
+- Save in format ready for Style-SFT and DPO
+"""
+
+import json
+import random
+import os
+from pathlib import Path
+
+def load_raw_data(path: str) -> list:
+    data = []
+    with open(path) as f:
+        for line in f:
+            d = json.loads(line)
+            data.append(d)
+    return data
+
+def filter_data(data: list, min_words: int = 50, max_words: int = 800) -> list:
+    """Filter for quality samples."""
+    filtered = []
+    for d in data:
+        hw = d['human_words']
+        aw = d['ai_words']
+        # Both texts should be within reasonable range
+        if hw < min_words or aw < min_words:
+            continue
+        if hw > max_words or aw > max_words:
+            continue
+        # Length ratio should be reasonable (AI version shouldn't be too different)
+        ratio = aw / hw if hw > 0 else 0
+        if ratio < 0.5 or ratio > 2.0:
+            continue
+        # Text should not be empty or too short
+        if len(d['human_text'].strip()) < 100 or len(d['ai_text'].strip()) < 100:
+            continue
+        filtered.append(d)
+    return filtered
+
+def prepare_sft_data(data: list) -> list:
+    """
+    Prepare data for Style-injection SFT.
+    Each sample has:
+    - input: AI text
+    - target_human: human text (for style transfer task)
+    - target_ai: AI text (for reconstruction task)
+    """
+    sft_data = []
+    for d in data:
+        sft_data.append({
+            'id': d['essay_id'],
+            'type': d['type'],
+            'tier': d.get('tier', 'unknown'),
+            'input_text': d['ai_text'],
+            'human_text': d['human_text'],
+            'ai_text': d['ai_text'],
+        })
+    return sft_data
+
+def split_data(data: list, val_ratio: float = 0.1, seed: int = 42) -> tuple:
+    """Split into train and validation sets, stratified by type."""
+    random.seed(seed)
+    
+    # Separate by type
+    ps_data = [d for d in data if d['type'] == 'ps']
+    supp_data = [d for d in data if d['type'] == 'supp']
+    
+    random.shuffle(ps_data)
+    random.shuffle(supp_data)
+    
+    ps_val_size = max(1, int(len(ps_data) * val_ratio))
+    supp_val_size = max(1, int(len(supp_data) * val_ratio))
+    
+    val_data = ps_data[:ps_val_size] + supp_data[:supp_val_size]
+    train_data = ps_data[ps_val_size:] + supp_data[supp_val_size:]
+    
+    random.shuffle(train_data)
+    random.shuffle(val_data)
+    
+    return train_data, val_data
+
+def save_jsonl(data: list, path: str):
+    with open(path, 'w') as f:
+        for d in data:
+            f.write(json.dumps(d, ensure_ascii=False) + '\n')
+
+def main():
+    raw_path = '/home/ubuntu/experiment/training_pairs_v3_final.jsonl'
+    output_dir = '/home/ubuntu/mash_training/data'
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Load and filter
+    print("Loading raw data...")
+    raw_data = load_raw_data(raw_path)
+    print(f"  Raw samples: {len(raw_data)}")
+    
+    print("Filtering data...")
+    filtered = filter_data(raw_data)
+    print(f"  After filtering: {len(filtered)}")
+    
+    # Prepare SFT format
+    print("Preparing SFT data...")
+    sft_data = prepare_sft_data(filtered)
+    
+    # Split
+    print("Splitting into train/val...")
+    train_data, val_data = split_data(sft_data)
+    print(f"  Train: {len(train_data)}")
+    print(f"  Val: {len(val_data)}")
+    
+    # Type distribution
+    from collections import Counter
+    train_types = Counter(d['type'] for d in train_data)
+    val_types = Counter(d['type'] for d in val_data)
+    print(f"  Train types: {dict(train_types)}")
+    print(f"  Val types: {dict(val_types)}")
+    
+    # Save
+    save_jsonl(train_data, os.path.join(output_dir, 'train.jsonl'))
+    save_jsonl(val_data, os.path.join(output_dir, 'val.jsonl'))
+    save_jsonl(sft_data, os.path.join(output_dir, 'all.jsonl'))
+    
+    print(f"\nData saved to {output_dir}/")
+    print("  train.jsonl - for training")
+    print("  val.jsonl   - for validation")
+    print("  all.jsonl   - complete dataset")
+
+if __name__ == '__main__':
+    main()