Upload plm.py with huggingface_hub

plm.py

@@ -0,0 +1,622 @@
+# -*- coding: utf-8 -*-
+"""
+Progressive LoRA Merging (PLM)
+Complete model identity replacement via iterative train-merge cycles.
+
+Paper: "Body Snatching: Complete Model Identity Replacement via Progressive LoRA Merging"
+Author: Ouissam Said Drissi (wissam.idrissi@gmail.com)
+
+Usage:
+    python plm.py --base-model Qwen/Qwen3-1.7B --dataset your_data.jsonl --cycles 100
+    python plm.py --base-model meta-llama/Llama-3-8B --dataset data.jsonl --cycles 50
+
+The key insight: Catastrophic forgetting is a FEATURE, not a bug.
+Each cycle permanently merges learned weights into the base, progressively
+replacing the model's original identity with your data.
+"""
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TrainingArguments,
+    Trainer,
+    TrainerCallback,
+    BitsAndBytesConfig,
+)
+from peft import LoraConfig, get_peft_model, PeftModel, prepare_model_for_kbit_training
+from dataclasses import dataclass
+from typing import Dict, List, Any, Optional
+from datasets import Dataset
+import json
+import pandas as pd
+from tqdm import tqdm
+import random
+import shutil
+from pathlib import Path
+import gc
+import argparse
+import os
+from datetime import datetime
+
+
+# =============================================================================
+# CONFIGURATION
+# =============================================================================
+
+DEFAULT_CONFIG = {
+    "lora_r": 8,                    # LoRA rank (small is fine, we accumulate over cycles)
+    "lora_alpha": 32,               # LoRA alpha (4:1 ratio with rank)
+    "lora_dropout": 0.05,           # Light dropout
+    "learning_rate": 1e-4,          # Standard LoRA learning rate
+    "epochs_per_cycle": 1,          # Epochs before each merge
+    "batch_size": 1,                # Per-device batch size
+    "gradient_accumulation": 4,     # Effective batch = batch_size * this
+    "max_length": 4096,             # Max sequence length
+    "warmup_steps": 50,             # Warmup steps per cycle
+    "save_every_n_cycles": 5,       # Save checkpoint every N cycles
+    "output_dir": "./plm_output",   # Output directory
+}
+
+
+# =============================================================================
+# DATA LOADING
+# =============================================================================
+
+def load_dataset_jsonl(file_path: str, tokenizer, max_length: int = 4096) -> List[str]:
+    """
+    Load dataset from JSONL file.
+    
+    Expected format (any of these):
+        {"text": "full conversation text"}
+        {"prompt": "...", "response": "..."}
+        {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
+    """
+    print(f"\nLoading dataset from {file_path}...")
+    
+    texts = []
+    skipped = 0
+    
+    with open(file_path, 'r', encoding='utf-8') as f:
+        for line_num, line in enumerate(f, 1):
+            if not line.strip():
+                continue
+            
+            try:
+                data = json.loads(line)
+            except json.JSONDecodeError as e:
+                print(f"  [Skip] Line {line_num}: Invalid JSON - {str(e)[:50]}")
+                skipped += 1
+                continue
+            
+            # Handle different formats
+            if 'text' in data:
+                text = data['text']
+            elif 'training_data' in data:
+                text = data['training_data']
+            elif 'prompt' in data and 'response' in data:
+                # Convert to chat format
+                text = f"<|im_start|>user\n{data['prompt']}<|im_end|>\n<|im_start|>assistant\n{data['response']}<|im_end|>"
+            elif 'messages' in data:
+                # Convert messages array to text
+                text = ""
+                for msg in data['messages']:
+                    role = msg.get('role', 'user')
+                    content = msg.get('content', '')
+                    text += f"<|im_start|>{role}\n{content}<|im_end|>\n"
+                text = text.strip()
+            else:
+                print(f"  [Skip] Line {line_num}: Unknown format - {list(data.keys())}")
+                skipped += 1
+                continue
+            
+            # Check length
+            token_count = len(tokenizer.encode(text, add_special_tokens=False))
+            if token_count > max_length:
+                skipped += 1
+                continue
+            
+            texts.append(text)
+    
+    print(f"  Loaded: {len(texts)} examples")
+    if skipped > 0:
+        print(f"  Skipped: {skipped} examples")
+    
+    random.shuffle(texts)
+    return texts
+
+
+# =============================================================================
+# MODEL LOADING
+# =============================================================================
+
+def load_model_4bit(model_path: str):
+    """Load model in 4-bit quantization for memory-efficient training."""
+    
+    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+    dtype = torch.bfloat16 if use_bf16 else torch.float16
+    
+    print(f"\n=== Loading Model (4-bit) ===")
+    print(f"Model: {model_path}")
+    print(f"Compute dtype: {'BF16' if use_bf16 else 'FP16'}")
+    
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=dtype,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        torch_dtype=dtype,
+        device_map={"": 0},
+        trust_remote_code=True,
+        use_cache=False,
+        low_cpu_mem_usage=True,
+        quantization_config=bnb_config,
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path,
+        trust_remote_code=True,
+        padding_side="right"
+    )
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model.config.pad_token_id = tokenizer.pad_token_id
+    
+    print(f"  Loaded successfully")
+    print(f"  Vocab size: {len(tokenizer)}")
+    
+    return model, tokenizer
+
+
+def load_model_full_precision(model_path: str, tokenizer):
+    """Load model in full precision (BF16) for merging."""
+    
+    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+    dtype = torch.bfloat16 if use_bf16 else torch.float16
+    
+    print(f"\n=== Loading Model (Full Precision for Merge) ===")
+    print(f"Model: {model_path}")
+    print(f"Dtype: {dtype}")
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        torch_dtype=dtype,
+        device_map="cpu",  # CPU for merge to save VRAM
+        trust_remote_code=True,
+        low_cpu_mem_usage=True,
+    )
+    
+    # Resize embeddings to match tokenizer
+    model.resize_token_embeddings(len(tokenizer))
+    
+    return model
+
+
+# =============================================================================
+# LORA SETUP
+# =============================================================================
+
+def apply_lora(model, config: dict):
+    """Apply fresh LoRA adapters to model."""
+    
+    print(f"\n=== Applying LoRA ===")
+    print(f"  Rank: {config['lora_r']}, Alpha: {config['lora_alpha']}")
+    
+    # Prepare for k-bit training
+    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
+    
+    lora_config = LoraConfig(
+        r=config['lora_r'],
+        lora_alpha=config['lora_alpha'],
+        lora_dropout=config['lora_dropout'],
+        target_modules="all-linear",
+        bias="none",
+        task_type="CAUSAL_LM"
+    )
+    
+    model = get_peft_model(model, lora_config)
+    
+    # Print stats
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    print(f"  Trainable: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)")
+    
+    return model
+
+
+# =============================================================================
+# MERGING
+# =============================================================================
+
+def merge_lora_high_precision(adapter_path: str, base_model_path: str, output_path: str, tokenizer):
+    """
+    Merge LoRA adapter into base model using high precision (BF16).
+    
+    CRITICAL: Always merge in full precision, never in 4-bit!
+    """
+    print(f"\n=== Merging LoRA (High Precision) ===")
+    print(f"  Base: {base_model_path}")
+    print(f"  Adapter: {adapter_path}")
+    print(f"  Output: {output_path}")
+    
+    # Load base in full precision
+    base_model = load_model_full_precision(base_model_path, tokenizer)
+    
+    # Apply adapter
+    print("  Applying adapter...")
+    model = PeftModel.from_pretrained(base_model, adapter_path)
+    
+    # Merge
+    print("  Merging weights...")
+    merged = model.merge_and_unload()
+    
+    # Save
+    output_dir = Path(output_path)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    merged.save_pretrained(output_dir, safe_serialization=True)
+    tokenizer.save_pretrained(output_dir)
+    
+    print(f"  Saved to: {output_dir}")
+    
+    # Cleanup
+    del merged, model, base_model
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    
+    return str(output_dir)
+
+
+# =============================================================================
+# TOKENIZATION
+# =============================================================================
+
+def tokenize_for_training(examples: dict, tokenizer, max_length: int) -> dict:
+    """Tokenize with causal LM labels."""
+    
+    encodings = tokenizer(
+        examples["text"],
+        max_length=max_length,
+        padding=False,
+        truncation=True,
+        return_tensors=None,
+    )
+    
+    # For causal LM, labels = input_ids
+    encodings["labels"] = encodings["input_ids"].copy()
+    
+    return encodings
+
+
+@dataclass
+class DataCollator:
+    """Collator that handles padding."""
+    tokenizer: Any
+    
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        input_ids = [torch.tensor(f["input_ids"]) for f in features]
+        labels = [torch.tensor(f["labels"]) for f in features]
+        
+        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id)
+        labels = pad_sequence(labels, batch_first=True, padding_value=-100)
+        attention_mask = (input_ids != self.tokenizer.pad_token_id).long()
+        
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels
+        }
+
+
+# =============================================================================
+# TRAINING
+# =============================================================================
+
+class ProgressCallback(TrainerCallback):
+    """Simple progress tracking."""
+    
+    def __init__(self, cycle: int):
+        self.cycle = cycle
+        self.losses = []
+    
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs and 'loss' in logs:
+            self.losses.append(logs['loss'])
+            avg = sum(self.losses[-50:]) / min(50, len(self.losses))
+            print(f"\r  [Cycle {self.cycle}] Step {state.global_step} | Loss: {logs['loss']:.4f} | Avg: {avg:.4f}", end="")
+
+
+def train_one_cycle(model, tokenizer, texts: List[str], cycle: int, config: dict):
+    """Train for one cycle (one or more epochs)."""
+    
+    print(f"\n{'='*60}")
+    print(f"CYCLE {cycle}")
+    print(f"{'='*60}")
+    print(f"  Examples: {len(texts)}")
+    
+    # Create dataset
+    df = pd.DataFrame({"text": texts})
+    train_size = int(0.95 * len(df))
+    
+    train_dataset = Dataset.from_pandas(df[:train_size])
+    eval_dataset = Dataset.from_pandas(df[train_size:])
+    
+    # Tokenize
+    train_dataset = train_dataset.map(
+        lambda x: tokenize_for_training(x, tokenizer, config['max_length']),
+        batched=True,
+        remove_columns=train_dataset.column_names,
+    )
+    eval_dataset = eval_dataset.map(
+        lambda x: tokenize_for_training(x, tokenizer, config['max_length']),
+        batched=True,
+        remove_columns=eval_dataset.column_names,
+    )
+    
+    # Training args
+    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+    
+    training_args = TrainingArguments(
+        output_dir=f"{config['output_dir']}/cycle_{cycle}",
+        num_train_epochs=config['epochs_per_cycle'],
+        per_device_train_batch_size=config['batch_size'],
+        per_device_eval_batch_size=config['batch_size'],
+        gradient_accumulation_steps=config['gradient_accumulation'],
+        warmup_steps=config['warmup_steps'],
+        learning_rate=config['learning_rate'],
+        bf16=use_bf16,
+        fp16=not use_bf16,
+        logging_steps=10,
+        eval_strategy="epoch",
+        save_strategy="no",
+        report_to="none",
+        disable_tqdm=True,
+        gradient_checkpointing=True,
+    )
+    
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        processing_class=tokenizer,
+        data_collator=DataCollator(tokenizer),
+        callbacks=[ProgressCallback(cycle)],
+    )
+    
+    # Train
+    trainer.train()
+    print()  # Newline after progress
+    
+    # Get final loss
+    eval_results = trainer.evaluate()
+    print(f"  Eval Loss: {eval_results['eval_loss']:.4f}")
+    
+    return model, eval_results['eval_loss']
+
+
+# =============================================================================
+# MAIN PROGRESSIVE LOOP
+# =============================================================================
+
+def progressive_lora_merge(
+    base_model: str,
+    dataset_path: str,
+    num_cycles: int,
+    config: dict = None
+) -> str:
+    """
+    Main Progressive LoRA Merging loop.
+    
+    For each cycle:
+    1. Load base model (4-bit for training)
+    2. Apply fresh LoRA
+    3. Train
+    4. Save adapter
+    5. Merge in high precision (BF16)
+    6. Use merged as new base
+    7. Repeat
+    
+    Returns path to final merged model.
+    """
+    
+    if config is None:
+        config = DEFAULT_CONFIG.copy()
+    
+    output_dir = Path(config['output_dir'])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    print("\n" + "="*60)
+    print("PROGRESSIVE LORA MERGING")
+    print("="*60)
+    print(f"Base Model: {base_model}")
+    print(f"Dataset: {dataset_path}")
+    print(f"Cycles: {num_cycles}")
+    print(f"Output: {output_dir}")
+    print("="*60)
+    
+    # Track state
+    current_base = base_model
+    best_loss = float('inf')
+    best_cycle = 0
+    
+    # Initial model load to get tokenizer
+    model, tokenizer = load_model_4bit(base_model)
+    
+    # Load dataset
+    texts = load_dataset_jsonl(dataset_path, tokenizer, config['max_length'])
+    if len(texts) == 0:
+        raise ValueError("No valid examples in dataset!")
+    
+    # Save config
+    with open(output_dir / "config.json", 'w') as f:
+        json.dump({
+            "base_model": base_model,
+            "dataset": dataset_path,
+            "num_cycles": num_cycles,
+            "config": config,
+            "started": datetime.now().isoformat()
+        }, f, indent=2)
+    
+    # Main loop
+    for cycle in range(1, num_cycles + 1):
+        
+        # Apply fresh LoRA
+        if cycle == 1:
+            model = apply_lora(model, config)
+        else:
+            # Reload from merged base
+            del model
+            torch.cuda.empty_cache()
+            gc.collect()
+            
+            model, tokenizer = load_model_4bit(current_base)
+            model = apply_lora(model, config)
+        
+        # Train
+        random.shuffle(texts)  # Reshuffle each cycle
+        model, eval_loss = train_one_cycle(model, tokenizer, texts, cycle, config)
+        
+        # Track best
+        if eval_loss < best_loss:
+            best_loss = eval_loss
+            best_cycle = cycle
+            print(f"  ★ New best loss!")
+        
+        # Save adapter
+        adapter_path = output_dir / f"adapters/cycle_{cycle}"
+        adapter_path.mkdir(parents=True, exist_ok=True)
+        model.save_pretrained(adapter_path)
+        tokenizer.save_pretrained(adapter_path)
+        
+        # Merge
+        merged_path = output_dir / f"merged/cycle_{cycle}"
+        
+        del model
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        merge_lora_high_precision(
+            str(adapter_path),
+            current_base,
+            str(merged_path),
+            tokenizer
+        )
+        
+        # Update base for next cycle
+        current_base = str(merged_path)
+        
+        # Periodic checkpoint
+        if cycle % config['save_every_n_cycles'] == 0:
+            checkpoint_path = output_dir / "checkpoints" / f"cycle_{cycle}"
+            shutil.copytree(merged_path, checkpoint_path, dirs_exist_ok=True)
+            print(f"  Checkpoint saved: {checkpoint_path}")
+        
+        # Cleanup old merged (keep disk space manageable)
+        if cycle > 1:
+            old_merged = output_dir / f"merged/cycle_{cycle-1}"
+            if old_merged.exists() and cycle % config['save_every_n_cycles'] != 1:
+                shutil.rmtree(old_merged)
+        
+        print(f"  Cycle {cycle} complete. New base: {current_base}")
+    
+    # Final save
+    final_path = output_dir / "final_model"
+    shutil.copytree(current_base, final_path, dirs_exist_ok=True)
+    
+    # Summary
+    print("\n" + "="*60)
+    print("TRAINING COMPLETE")
+    print("="*60)
+    print(f"Total cycles: {num_cycles}")
+    print(f"Best loss: {best_loss:.4f} (cycle {best_cycle})")
+    print(f"Final model: {final_path}")
+    print("="*60)
+    
+    # Save final state
+    with open(output_dir / "results.json", 'w') as f:
+        json.dump({
+            "total_cycles": num_cycles,
+            "best_loss": best_loss,
+            "best_cycle": best_cycle,
+            "final_model": str(final_path),
+            "completed": datetime.now().isoformat()
+        }, f, indent=2)
+    
+    return str(final_path)
+
+
+# =============================================================================
+# CLI
+# =============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Progressive LoRA Merging - Complete model identity replacement",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python plm.py --base-model Qwen/Qwen3-1.7B --dataset data.jsonl --cycles 100
+  python plm.py --base-model meta-llama/Llama-3-8B --dataset data.jsonl --cycles 50 --lora-r 16
+
+Dataset format (JSONL, any of these):
+  {"text": "full conversation text"}
+  {"prompt": "user input", "response": "assistant output"}
+  {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
+
+Paper: "Body Snatching: Complete Model Identity Replacement via Progressive LoRA Merging"
+Author: Ouissam Said Drissi (wissam.idrissi@gmail.com)
+        """
+    )
+    
+    # Required
+    parser.add_argument("--base-model", required=True, help="Base model path or HF model ID")
+    parser.add_argument("--dataset", required=True, help="Path to JSONL dataset")
+    parser.add_argument("--cycles", type=int, required=True, help="Number of train-merge cycles")
+    
+    # Optional
+    parser.add_argument("--output-dir", default="./plm_output", help="Output directory")
+    parser.add_argument("--lora-r", type=int, default=8, help="LoRA rank")
+    parser.add_argument("--lora-alpha", type=int, default=32, help="LoRA alpha")
+    parser.add_argument("--learning-rate", type=float, default=1e-4, help="Learning rate")
+    parser.add_argument("--batch-size", type=int, default=1, help="Batch size")
+    parser.add_argument("--max-length", type=int, default=4096, help="Max sequence length")
+    parser.add_argument("--epochs-per-cycle", type=int, default=1, help="Epochs per cycle")
+    parser.add_argument("--save-every", type=int, default=5, help="Save checkpoint every N cycles")
+    
+    args = parser.parse_args()
+    
+    # Build config
+    config = DEFAULT_CONFIG.copy()
+    config.update({
+        "output_dir": args.output_dir,
+        "lora_r": args.lora_r,
+        "lora_alpha": args.lora_alpha,
+        "learning_rate": args.learning_rate,
+        "batch_size": args.batch_size,
+        "max_length": args.max_length,
+        "epochs_per_cycle": args.epochs_per_cycle,
+        "save_every_n_cycles": args.save_every,
+    })
+    
+    # Run
+    final_model = progressive_lora_merge(
+        base_model=args.base_model,
+        dataset_path=args.dataset,
+        num_cycles=args.cycles,
+        config=config
+    )
+    
+    print(f"\nDone! Final model at: {final_model}")
+
+
+if __name__ == "__main__":
+    main()