train.py

import torch, gc, os, numpy as np, evaluate, json
from datasets import load_dataset
from transformers import (
    AutoTokenizer, AutoModelForQuestionAnswering,
    TrainingArguments, Trainer, default_data_collator
)
from peft import LoraConfig, get_peft_model, TaskType
from huggingface_hub import login
import sys

def main():
    # Get model name from environment
    model_name = os.environ.get('MODEL_NAME', 'roberta-cuad-qa')
    
    # Login to HF Hub
    hf_token = os.environ.get('roberta_token')
    if hf_token:
        login(token=hf_token)
        print("✅ Logged into Hugging Face Hub")
    else:
        print("⚠️ No HF_TOKEN found - model won't be pushed to Hub")
    
    # Setup
    torch.cuda.empty_cache()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"🔧 Using device: {device}")
    
    if torch.cuda.is_available():
        print(f"🎯 GPU: {torch.cuda.get_device_name()}")
        print(f"💾 GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    
    # Load and prepare data
    print("📚 Loading CUAD dataset...")
    raw = load_dataset("theatticusproject/cuad-qa", split="train", trust_remote_code=True)
    
    # Use subset for faster training on free GPU
    N = 2000  
    raw = raw.shuffle(seed=42).select(range(min(N, len(raw))))
    ds = raw.train_test_split(test_size=0.1, seed=42)
    train_ds, val_ds = ds["train"], ds["test"]
    
    print(f"✅ Data loaded - Train: {len(train_ds)}, Val: {len(val_ds)}")
    
    # Store original validation data for metrics
    original_val_data = [ex["answers"] for ex in val_ds]
    
    # Load model and tokenizer
    print("🤖 Loading RoBERTa model...")
    base_model = "roberta-base"
    tok = AutoTokenizer.from_pretrained(base_model, use_fast=True)
    model = AutoModelForQuestionAnswering.from_pretrained(base_model)
    
    # Add LoRA
    print("🔧 Adding LoRA adapters...")
    lora_cfg = LoraConfig(
        task_type=TaskType.QUESTION_ANS,
        target_modules=["query", "value"],
        r=16,
        lora_alpha=32,  # Improved scaling
        lora_dropout=0.05,
    )
    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()
    model.to(device)
    
    # Tokenization function
    max_len, doc_stride = 384, 128
    
    def preprocess(examples):
        tok_batch = tok(
            examples["question"],
            examples["context"],
            truncation="only_second",
            max_length=max_len,
            stride=doc_stride,
            return_overflowing_tokens=True,
            return_offsets_mapping=True,
            padding="max_length",
        )

        sample_map = tok_batch.pop("overflow_to_sample_mapping")
        offset_map = tok_batch.pop("offset_mapping")
        start_pos, end_pos = [], []

        for i, offsets in enumerate(offset_map):
            cls_idx = tok_batch["input_ids"][i].index(tok.cls_token_id)
            sample_idx = sample_map[i]
            answer = examples["answers"][sample_idx]

            if len(answer["answer_start"]) == 0:
                start_pos.append(cls_idx)
                end_pos.append(cls_idx)
                continue

            s_char = answer["answer_start"][0]
            e_char = s_char + len(answer["text"][0])
            seq_ids = tok_batch.sequence_ids(i)
            c0, c1 = seq_ids.index(1), len(seq_ids) - 1 - seq_ids[::-1].index(1)

            if not (offsets[c0][0] <= s_char <= offsets[c1][1]):
                start_pos.append(cls_idx)
                end_pos.append(cls_idx)
                continue

            st = c0
            while st <= c1 and offsets[st][0] <= s_char: 
                st += 1
            en = c1
            while en >= c0 and offsets[en][1] >= e_char: 
                en -= 1

            # Fixed position calculation with bounds checking
            start_pos.append(max(c0, min(st - 1, c1)))
            end_pos.append(max(c0, min(en + 1, c1)))

        tok_batch["start_positions"] = start_pos
        tok_batch["end_positions"] = end_pos
        return tok_batch

    # Tokenize datasets
    print("🔄 Tokenizing datasets...")
    train_tok = train_ds.map(
        preprocess, batched=True, batch_size=100,
        remove_columns=train_ds.column_names, 
        desc="Tokenizing train"
    )
    val_tok = val_ds.map(
        preprocess, batched=True, batch_size=100,
        remove_columns=val_ds.column_names, 
        desc="Tokenizing validation"
    )
    
    # Clean up memory
    del raw, ds, train_ds, val_ds
    gc.collect()
    torch.cuda.empty_cache()
    
    # Metrics setup
    metric = evaluate.load("squad")

    def postprocess(preds, dataset):
        starts, ends = preds
        answers = []
        for i in range(len(starts)):
            a, b = int(np.argmax(starts[i])), int(np.argmax(ends[i]))
            if a > b: a, b = b, a
            text = tok.decode(dataset[i]["input_ids"][a:b+1], skip_special_tokens=True)
            answers.append(text.strip())
        return answers

    def compute_metrics(eval_pred):
        try:
            preds, _ = eval_pred
            texts = postprocess(preds, val_tok)
            predictions = [{"id": str(i), "prediction_text": t} for i, t in enumerate(texts)]
            references = [{"id": str(i), "answers": ans} for i, ans in enumerate(original_val_data)]
            return metric.compute(predictions=predictions, references=references)
        except Exception as e:
            print(f"⚠️ Metrics computation failed: {e}")
            return {"exact_match": 0.0, "f1": 0.0}

    # Training arguments
    output_dir = "./model_output"
    args = TrainingArguments(
        output_dir=output_dir,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=4,
        gradient_accumulation_steps=8,
        num_train_epochs=2,
        learning_rate=5e-4,
        lr_scheduler_type="cosine",
        warmup_ratio=0.1,
        fp16=True,
        eval_strategy="steps",
        eval_steps=250,
        save_steps=500,
        save_total_limit=2,
        logging_steps=50,
        weight_decay=0.01,
        remove_unused_columns=True,
        report_to=None,
        push_to_hub=False,  # We'll do this manually
        dataloader_pin_memory=False,  # Save memory
    )

    # Create trainer
    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_tok,
        eval_dataset=val_tok,
        tokenizer=tok,
        data_collator=default_data_collator,
        compute_metrics=compute_metrics,
    )

    print(f"🚀 Starting training...")
    print(f"📊 Total training samples: {len(train_tok)}")
    print(f"📊 Total validation samples: {len(val_tok)}")
    
    if torch.cuda.is_available():
        print(f"💾 GPU memory before training: {torch.cuda.memory_allocated()/1024**3:.2f} GB")

    # Training loop with error handling
    try:
        trainer.train()
        print("✅ Training completed successfully!")
        
    except RuntimeError as e:
        if "CUDA out of memory" in str(e):
            print("⚠️ GPU OOM - reducing batch size and retrying...")
            torch.cuda.empty_cache()
            gc.collect()
            
            # Reduce batch size
            args.per_device_train_batch_size = 1
            args.gradient_accumulation_steps = 16
            
            trainer = Trainer(
                model=model, args=args,
                train_dataset=train_tok, eval_dataset=val_tok,
                tokenizer=tok, data_collator=default_data_collator,
                compute_metrics=compute_metrics,
            )
            trainer.train()
            print("✅ Training completed with reduced batch size!")
        else:
            raise e

    # Save model locally first
    print("💾 Saving model locally...")
    os.makedirs(output_dir, exist_ok=True)
    trainer.model.save_pretrained(output_dir)
    tok.save_pretrained(output_dir)
    
    # Save training info
    training_info = {
        "model_name": model_name,
        "base_model": base_model,
        "dataset": "theatticusproject/cuad-qa",
        "lora_config": {
            "r": lora_cfg.r,
            "lora_alpha": lora_cfg.lora_alpha,
            "target_modules": lora_cfg.target_modules,
            "lora_dropout": lora_cfg.lora_dropout,
        },
        "training_samples": len(train_tok),
        "validation_samples": len(val_tok),
    }
    
    with open(os.path.join(output_dir, "training_info.json"), "w") as f:
        json.dump(training_info, f, indent=2)
    
    # Push to Hub if token available
    if hf_token:
        try:
            print(f"⬆️ Pushing model to Hub: {model_name}")
            trainer.model.push_to_hub(model_name, private=False)
            tok.push_to_hub(model_name, private=False)
            
            # Also push training info
            from huggingface_hub import upload_file
            upload_file(
                path_or_fileobj=os.path.join(output_dir, "training_info.json"),
                path_in_repo="training_info.json",
                repo_id=model_name,
                repo_type="model"
            )
            
            print(f"🎉 Model successfully saved to: https://huggingface.co/{model_name}")
            
        except Exception as e:
            print(f"❌ Failed to push to Hub: {e}")
            print("💾 Model saved locally in ./model_output/")
    else:
        print("💾 Model saved locally in ./model_output/ (no HF token for Hub upload)")

    print("🏁 Training pipeline completed!")

if __name__ == "__main__":
    main()