Create train1.py

train1.py

@@ -0,0 +1,379 @@
+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 - REDUCED SIZE FOR FASTER TRAINING
+    print("📚 Loading CUAD dataset...")
+    raw = load_dataset("theatticusproject/cuad-qa", split="train", trust_remote_code=True)
+    
+    # Use 5000 samples for good model quality - expect ~1 hour training
+    N = 5000  
+    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 
+    print("📊 Preparing metrics data...")
+    original_val_data = []
+    val_sample_mapping = []  # Track which tokenized sample maps to which original
+    
+    for i, ex in enumerate(val_ds):
+        original_val_data.append(ex["answers"])
+    
+    # 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,
+        lora_dropout=0.05,
+    )
+    model = get_peft_model(model, lora_cfg)
+    model.print_trainable_parameters()
+    model.to(device)
+    
+    # Tokenization function - AGGRESSIVE OPTIMIZATION TO PREVENT EXPANSION
+    max_len, doc_stride = 512, 400  # Much larger stride to minimize chunks per document
+    
+    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
+        
+        # Store sample mapping for metrics calculation
+        tok_batch["sample_mapping"] = sample_map
+        return tok_batch
+    
+    # Tokenize datasets
+    print("🔄 Tokenizing datasets...")
+    train_tok = train_ds.map(
+        preprocess,
+        batched=True,
+        batch_size=50,  # Smaller batch size for preprocessing
+        remove_columns=train_ds.column_names,
+        desc="Tokenizing train"
+    )
+    
+    val_tok = val_ds.map(
+        preprocess,
+        batched=True,
+        batch_size=50,
+        remove_columns=val_ds.column_names,
+        desc="Tokenizing validation"
+    )
+    
+    # DEBUG: Print actual dataset sizes after tokenization
+    print(f"🔍 DEBUG INFO:")
+    print(f"  Original samples: {N}")
+    print(f"  After tokenization - Train: {len(train_tok)}, Val: {len(val_tok)}")
+    print(f"  Expansion factor: {len(train_tok)/len(train_ds):.1f}x")
+    
+    # SAFETY CHECK: If expansion is too high, reduce data size automatically
+    expansion_factor = len(train_tok) / len(train_ds)
+    if expansion_factor > 12:  # Slightly more permissive for 4K samples
+        print(f"⚠️ HIGH EXPANSION DETECTED ({expansion_factor:.1f}x)!")
+        print("🔧 Auto-reducing dataset size to prevent excessively slow training...")
+        
+        # Allow up to 20k samples for 1 hour training
+        target_size = min(20000, len(train_tok))  # Max 20k samples
+        train_indices = list(range(0, len(train_tok), max(1, len(train_tok) // target_size)))[:target_size]
+        val_indices = list(range(0, len(val_tok), max(1, len(val_tok) // (target_size // 10))))[:target_size // 10]
+        
+        train_tok = train_tok.select(train_indices)
+        val_tok = val_tok.select(val_indices)
+        
+        print(f"✅ Reduced to - Train: {len(train_tok)}, Val: {len(val_tok)}")
+        print(f"📈 This should complete in ~45-75 minutes")
+    
+    # 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
+            starts, ends = preds
+            
+            # Group predictions by original sample (handle multiple chunks per sample)
+            sample_predictions = {}
+            for i in range(len(starts)):
+                # Get which original sample this tokenized example came from
+                if hasattr(val_tok[i], 'sample_mapping') and 'sample_mapping' in val_tok[i]:
+                    orig_idx = val_tok[i]['sample_mapping']
+                else:
+                    # Fallback: assume 1:1 mapping (may be inaccurate with chunking)
+                    orig_idx = min(i, len(original_val_data) - 1)
+                
+                # Get best answer span for this chunk
+                start_idx = int(np.argmax(starts[i]))
+                end_idx = int(np.argmax(ends[i]))
+                if start_idx > end_idx:
+                    start_idx, end_idx = end_idx, start_idx
+                
+                # Extract answer text
+                answer_text = tok.decode(
+                    val_tok[i]["input_ids"][start_idx:end_idx+1],
+                    skip_special_tokens=True
+                ).strip()
+                
+                # Store best prediction for this original sample
+                confidence = float(starts[i][start_idx]) + float(ends[i][end_idx])
+                if orig_idx not in sample_predictions or confidence > sample_predictions[orig_idx][1]:
+                    sample_predictions[orig_idx] = (answer_text, confidence)
+            
+            # Format for SQuAD metric
+            predictions = []
+            references = []
+            for orig_idx in range(len(original_val_data)):
+                pred_text = sample_predictions.get(orig_idx, ("", 0))[0]
+                predictions.append({
+                    "id": str(orig_idx),
+                    "prediction_text": pred_text
+                })
+                references.append({
+                    "id": str(orig_idx),
+                    "answers": original_val_data[orig_idx]
+                })
+            
+            result = metric.compute(predictions=predictions, references=references)
+            
+            # Add some debugging info
+            print(f"📊 Evaluation: EM={result['exact_match']:.3f}, F1={result['f1']:.3f}")
+            return result
+            
+        except Exception as e:
+            print(f"⚠️ Metrics computation failed: {e}")
+            print(f"  Pred shape: {np.array(preds).shape if preds else 'None'}")
+            print(f"  Val dataset size: {len(val_tok)}")
+            print(f"  Original val size: {len(original_val_data)}")
+            return {"exact_match": 0.0, "f1": 0.0}
+    
+    # OPTIMIZED Training arguments
+    output_dir = "./model_output"
+    args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=8,  # INCREASED from 2
+        per_device_eval_batch_size=8,   # INCREASED from 4
+        gradient_accumulation_steps=2,  # REDUCED from 8
+        num_train_epochs=3,             # Back to 3 epochs for better training
+        learning_rate=5e-4,
+        lr_scheduler_type="cosine",
+        warmup_ratio=0.1,
+        bf16=True,                      # CHANGED from fp16 (better for newer GPUs)
+        eval_strategy="steps",
+        eval_steps=100,                 # REDUCED from 250
+        save_steps=200,                 # REDUCED from 500
+        save_total_limit=2,
+        logging_steps=25,               # REDUCED from 50
+        weight_decay=0.01,
+        remove_unused_columns=True,
+        report_to=None,
+        push_to_hub=False,
+        dataloader_pin_memory=True,     # CHANGED to True for faster data loading
+        dataloader_num_workers=4,       # ADDED for parallel data loading
+        gradient_checkpointing=False,   # DISABLED to trade memory for speed
+    )
+    
+    # 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)}")
+    print(f"⚡ Effective batch size: {args.per_device_train_batch_size * args.gradient_accumulation_steps}")
+    
+    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 = 4
+            args.gradient_accumulation_steps = 4
+            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",
+        "original_samples": N,
+        "training_samples_after_tokenization": len(train_tok),
+        "validation_samples_after_tokenization": len(val_tok),
+        "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_args": {
+            "batch_size": args.per_device_train_batch_size,
+            "gradient_accumulation_steps": args.gradient_accumulation_steps,
+            "effective_batch_size": args.per_device_train_batch_size * args.gradient_accumulation_steps,
+            "epochs": args.num_train_epochs,
+            "learning_rate": args.learning_rate,
+        }
+    }
+    
+    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()