Create train.py

train.py

@@ -0,0 +1,281 @@
+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()