Update train.py

train.py

@@ -1,371 +1,206 @@
-#!/usr/bin/env python
-# train_cuad_lora_efficient.py - FIXED VERSION
-"""
-CUAD fine-tune with LoRA - Fixed for realistic training times
-"""
+import os
+import collections
+import string
+import re
+import numpy as np
 
-import os, json, random, gc, time
-from collections import defaultdict
-from pathlib import Path
-
-import torch, numpy as np
-from datasets import load_dataset, Dataset, disable_caching
+from datasets import load_dataset, load_metric
 from transformers import (
-    AutoTokenizer, AutoModelForQuestionAnswering,
-    TrainingArguments, default_data_collator, Trainer
+    DebertaTokenizerFast,
+    DebertaForQuestionAnswering,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
 )
-from peft import LoraConfig, get_peft_model, TaskType
-import evaluate
+from peft import LoraConfig, get_peft_model
 from huggingface_hub import login
 
-disable_caching()
-
-# Set tokenizers parallelism to avoid warnings
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-# ─────────────────────────────────────────────────────────────── config ──
-
-MAX_LEN = 512  # Slightly longer context
-DOC_STRIDE = 256   # Larger stride = fewer chunks = faster training
-SEED = 42
-BATCH_SIZE = 1000  # Process in larger, more efficient batches
-
-# Back to reasonable subset size since you've trained 5k before
-USE_SUBSET = True
-SUBSET_SIZE = 7000  # Good middle ground - more than your 5k success
-
-def set_seed(seed):
-    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-
-def balance_has_answer(dataset, ratio=2.0, max_samples=None):
-    """Keep all has-answer rows, down-sample no-answer rows to `ratio`."""
-    has, no = [], []
-    for ex in dataset:
-        (has if ex["answers"]["text"] else no).append(ex)
-    
-    print(f"📊 Original: {len(has)} has-answer, {len(no)} no-answer")
-    
-    # FIXED: Apply max_samples FIRST, then balance
-    if max_samples:
-        total_available = len(has) + len(no)
-        if total_available > max_samples:
-            # Sample proportionally from original distribution
-            has_ratio = len(has) / total_available
-            target_has = int(max_samples * has_ratio)
-            target_no = max_samples - target_has
-            
-            has = random.sample(has, min(target_has, len(has)))
-            no = random.sample(no, min(target_no, len(no)))
-            print(f"📉 Pre-balance subset: {len(has)} has-answer, {len(no)} no-answer")
-    
-    # Now balance within the subset
-    k = int(len(has) * ratio)
-    if len(no) > k:
-        no = random.sample(no, k)
-    
-    balanced = has + no
-    random.shuffle(balanced)  # Shuffle the final dataset
-    
-    print(f"📊 Final balanced: {len([x for x in balanced if x['answers']['text']])} has-answer, {len([x for x in balanced if not x['answers']['text']])} no-answer")
-    print(f"📊 Total examples: {len(balanced)}")
-    
-    return Dataset.from_list(balanced)
-
-# ────────────────────────────────────────────────────────────── postproc ──
-
-metric = evaluate.load("squad")
-
-def postprocess_qa(examples, features, raw_predictions, tokenizer):
-    """HF-style span extraction + n-best, returns SQuAD format dict."""
-    all_start, all_end = raw_predictions
-    example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
-    features_per_example = defaultdict(list)
-    for i, feat_id in enumerate(features["example_id"]):
-        features_per_example[example_id_to_index[feat_id]].append(i)
-
-    predictions = []
-
-    for example_idx, example in enumerate(examples):
-        best_score = -1e9
-        best_span = ""
-        context = example["context"]
-
-        for feat_idx in features_per_example[example_idx]:
-            start_logit = all_start[feat_idx]
-            end_logit = all_end[feat_idx]
-            offset = features["offset_mapping"][feat_idx]
-
-            start_idx = int(np.argmax(start_logit))
-            end_idx = int(np.argmax(end_logit))
-
-            if start_idx <= end_idx < len(offset):
-                start_char, _ = offset[start_idx]
-                _, end_char = offset[end_idx]
-                span = context[start_char:end_char].strip()
-                score = start_logit[start_idx] + end_logit[end_idx]
-                if score > best_score and span:
-                    best_score, best_span = score, span
-
-        predictions.append(
-            {"id": example["id"], "prediction_text": best_span}
-        )
-    return predictions
-
-# ───────────────────────────────────────────────────────────── preprocessing ──
-
-def preprocess_training_batch(examples, tokenizer):
-    """Training preprocessing - NO offset_mapping included"""
-    questions = examples["question"]
-    contexts = examples["context"]
-    
+# Load your HF token securely from environment variable
+hf_token = os.environ.get("roberta_token")
+if hf_token:
+    login(token=hf_token)
+else:
+    print("Warning: HF token not found in environment variable 'roberta_token'. Push to hub may fail.")
+
+metric = load_metric("squad")
+
+def normalize_answer(s):
+    """Lower text and remove punctuation/articles/extra whitespace"""
+    def remove_articles(text):
+        return re.sub(r'\b(a|an|the)\b', ' ', text)
+    def white_space_fix(text):
+        return ' '.join(text.split())
+    def remove_punc(text):
+        exclude = set(string.punctuation)
+        return ''.join(ch for ch in text if ch not in exclude)
+    def lower(text):
+        return text.lower()
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+def prepare_train_features(examples, tokenizer, max_length=512, doc_stride=128):
     tokenized_examples = tokenizer(
-        questions,
-        contexts,
+        examples["question"],
+        examples["context"],
         truncation="only_second",
-        max_length=MAX_LEN,
-        stride=DOC_STRIDE,
+        max_length=max_length,
+        stride=doc_stride,
         return_overflowing_tokens=True,
         return_offsets_mapping=True,
         padding="max_length",
     )
-    
     sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
     offset_mapping = tokenized_examples.pop("offset_mapping")
-    
+
     start_positions = []
     end_positions = []
-    
+
     for i, offsets in enumerate(offset_mapping):
-        cls_index = 0
+        input_ids = tokenized_examples["input_ids"][i]
+        cls_index = input_ids.index(tokenizer.cls_token_id)
         sample_index = sample_mapping[i]
         answers = examples["answers"][sample_index]
-        
-        if not answers["text"] or not answers["text"][0]:
+
+        if len(answers["answer_start"]) == 0:
             start_positions.append(cls_index)
             end_positions.append(cls_index)
-            continue
-        
-        answer_start_char = answers["answer_start"][0]
-        answer_text = answers["text"][0]
-        answer_end_char = answer_start_char + len(answer_text)
-        
-        token_start_index = cls_index
-        token_end_index = cls_index
-        
-        for token_index, (start_char, end_char) in enumerate(offsets):
-            if start_char <= answer_start_char < end_char:
-                token_start_index = token_index
-            if start_char < answer_end_char <= end_char:
-                token_end_index = token_index
-                break
-        
-        if token_start_index <= token_end_index and token_start_index > 0:
-            start_positions.append(token_start_index)
-            end_positions.append(token_end_index)
         else:
-            start_positions.append(cls_index)
-            end_positions.append(cls_index)
-    
+            start_char = answers["answer_start"][0]
+            end_char = start_char + len(answers["text"][0])
+            sequence_ids = tokenized_examples.sequence_ids(i)
+            token_start_index = 0
+            while sequence_ids[token_start_index] != 1:
+                token_start_index += 1
+            token_end_index = len(input_ids) - 1
+            while sequence_ids[token_end_index] != 1:
+                token_end_index -= 1
+            if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):
+                start_positions.append(cls_index)
+                end_positions.append(cls_index)
+            else:
+                while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
+                    token_start_index += 1
+                start_positions.append(token_start_index - 1)
+                while offsets[token_end_index][1] >= end_char:
+                    token_end_index -= 1
+                end_positions.append(token_end_index + 1)
+
     tokenized_examples["start_positions"] = start_positions
     tokenized_examples["end_positions"] = end_positions
-    
     return tokenized_examples
 
-def preprocess_validation_batch(examples, tokenizer):
-    """Validation preprocessing - INCLUDES offset_mapping and example_id"""
-    questions = examples["question"]
-    contexts = examples["context"]
-    
-    tokenized_examples = tokenizer(
-        questions,
-        contexts,
-        truncation="only_second",
-        max_length=MAX_LEN,
-        stride=DOC_STRIDE,
-        return_overflowing_tokens=True,
-        return_offsets_mapping=True,
-        padding="max_length",
-    )
-    
-    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
-    
-    tokenized_examples["example_id"] = [
-        examples["id"][sample_mapping[i]] for i in range(len(tokenized_examples["input_ids"]))
-    ]
-    
-    return tokenized_examples
-
-def preprocess_dataset_streaming(dataset, tokenizer, desc="Processing", is_training=True):
-    """Process dataset in batches using HuggingFace's map function with batching."""
-    print(f"🔄 {desc} dataset with batch processing...")
-    
-    if is_training:
-        preprocess_fn = preprocess_training_batch
-    else:
-        preprocess_fn = preprocess_validation_batch
-    
-    processed = dataset.map(
-        lambda examples: preprocess_fn(examples, tokenizer),
-        batched=True,
-        batch_size=BATCH_SIZE,
-        remove_columns=dataset.column_names,
-        desc=desc,
-        num_proc=1,
-    )
-    
-    print(f"✅ {desc} completed: {len(processed)} features")
-    return processed
+def postprocess_qa_predictions(examples, features, raw_predictions, n_best_size=20, max_answer_length=30):
+    all_start_logits, all_end_logits = raw_predictions
+    example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
+    features_per_example = collections.defaultdict(list)
+    for i, feature in enumerate(features):
+        features_per_example[example_id_to_index[feature["example_id"]]].append(i)
+
+    predictions = collections.OrderedDict()
+    for example_index, example in enumerate(examples):
+        feature_indices = features_per_example[example_index]
+        min_null_score = None
+        valid_answers = []
+        context = example["context"]
+        for feature_index in feature_indices:
+            start_logits = all_start_logits[feature_index]
+            end_logits = all_end_logits[feature_index]
+            offsets = features[feature_index]["offset_mapping"]
+            cls_index = features[feature_index]["input_ids"].index(features[feature_index]["cls_token_id"])
+            feature_null_score = start_logits[cls_index] + end_logits[cls_index]
+            if min_null_score is None or min_null_score > feature_null_score:
+                min_null_score = feature_null_score
+            start_indexes = np.argsort(start_logits)[-1: -n_best_size - 1: -1].tolist()
+            end_indexes = np.argsort(end_logits)[-1: -n_best_size - 1: -1].tolist()
+            for start_index in start_indexes:
+                for end_index in end_indexes:
+                    if (
+                        start_index >= len(offsets)
+                        or end_index >= len(offsets)
+                        or offsets[start_index] is None
+                        or offsets[end_index] is None
+                    ):
+                        continue
+                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:
+                        continue
+                    start_char = offsets[start_index][0]
+                    end_char = offsets[end_index][1]
+                    valid_answers.append(
+                        {"score": start_logits[start_index] + end_logits[end_index], "text": context[start_char:end_char]}
+                    )
+        best_answer = max(valid_answers, key=lambda x: x["score"]) if valid_answers else {"text": "", "score": 0.0}
+        predictions[example["id"]] = best_answer["text"]
+    return predictions
 
-# ───────────────────────────────────────────────────────────────── main ──
+def compute_metrics(p, tokenizer, examples, features):
+    predictions = postprocess_qa_predictions(examples, features, p.predictions)
+    formatted_predictions = [{"id": k, "prediction_text": v} for k, v in predictions.items()]
+    references = [{"id": ex["id"], "answers": ex["answers"]} for ex in examples]
+    return metric.compute(predictions=formatted_predictions, references=references)
 
 def main():
-    set_seed(SEED)
-
-    model_repo = os.getenv("MODEL_NAME", "AvocadoMuffin/roberta-cuad-qa-v4")
-
-    if (tokn := os.getenv("roberta_token")):
-        try:  
-            login(tokn)
-            print("🔑 HuggingFace Hub login OK")
-        except Exception as e: 
-            print(f"⚠️ Hub login failed: {e}")
-            tokn = None
-
-    print("📚 Loading CUAD…")
-    try:
-        cuad = load_dataset("theatticusproject/cuad-qa", split="train", trust_remote_code=True)
-        print(f"✅ Loaded {len(cuad)} examples")
-    except Exception as e:
-        print(f"❌ Dataset loading failed: {e}")
-        cuad = load_dataset("theatticusproject/cuad-qa", split="train", trust_remote_code=True, download_mode="force_redownload")
-    
-    cuad = cuad.shuffle(seed=SEED)
-    
-    # FIXED: Apply subset reduction more aggressively
-    subset_size = SUBSET_SIZE if USE_SUBSET else None
-    cuad = balance_has_answer(cuad, ratio=1.5, max_samples=subset_size)  # Reduced ratio too
-    print(f"📊 Final dataset size: {len(cuad)} examples")
-
-    # Estimate features after preprocessing
-    avg_features_per_example = 2.5  # Conservative estimate with stride
-    estimated_features = len(cuad) * avg_features_per_example
-    print(f"📊 Estimated training features: ~{int(estimated_features)}")
-
-    ds = cuad.train_test_split(test_size=0.1, seed=SEED)
-    train_raw, val_raw = ds["train"], ds["test"]
-
-    # ── tokeniser & model ──
-    base_ckpt = "deepset/roberta-base-squad2"
-    tok = AutoTokenizer.from_pretrained(base_ckpt, use_fast=True)
-    model = AutoModelForQuestionAnswering.from_pretrained(base_ckpt)
-
-    # FIXED: Lighter LoRA config for faster training
-    lora = LoraConfig(
-        task_type=TaskType.QUESTION_ANS,
-        r=16,  # Reduced from 32
-        lora_alpha=32,  # Reduced from 64
+    model_name = "microsoft/deberta-xlarge"
+    output_dir = "./deberta-lora-cuad-finetuned"
+
+    datasets = load_dataset("theatticusproject/cuad-qa")
+    tokenizer = DebertaTokenizerFast.from_pretrained(model_name)
+    model = DebertaForQuestionAnswering.from_pretrained(model_name)
+
+    # LoRA config: tune rank and dropout as needed
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=32,
+        target_modules=["query", "value"],  # Adjust for DeBERTa internals as needed
         lora_dropout=0.1,
-        target_modules=["query", "value"],  # Fewer modules
+        bias="none",
+        task_type="QUESTION_ANSWERING"
     )
-    model = get_peft_model(model, lora)
-    model.print_trainable_parameters()
-
-    # ── preprocessing ─────────────────────────────────────────
-    print("🔄 Starting preprocessing...")
-    
-    train_feats = preprocess_dataset_streaming(train_raw, tok, "Training", is_training=True)
-    val_feats = preprocess_dataset_streaming(val_raw, tok, "Validation", is_training=False)
+    model = get_peft_model(model, lora_config)
 
-    print(f"✅ Preprocessing completed!")
-    print(f"   Training features: {len(train_feats)}")
-    print(f"   Validation features: {len(val_feats)}")
+    train_dataset = datasets["train"].map(
+        lambda examples: prepare_train_features(examples, tokenizer),
+        batched=True,
+        remove_columns=datasets["train"].column_names,
+    )
+    val_dataset = datasets["validation"].map(
+        lambda examples: prepare_train_features(examples, tokenizer),
+        batched=True,
+        remove_columns=datasets["validation"].column_names,
+    )
 
-    # ── training args - FIXED for reasonable training time ──
-    batch_size = 16  # Good balance
-    gradient_accumulation_steps = 2
-    effective_batch_size = batch_size * gradient_accumulation_steps
-    
-    num_epochs = 3  # Keep it reasonable
-    steps_per_epoch = len(train_feats) // effective_batch_size
-    total_steps = steps_per_epoch * num_epochs
-    
-    eval_steps = max(25, steps_per_epoch // 8)  # More frequent eval
-    save_steps = eval_steps * 3
-    
-    print(f"📊 Training configuration:")
-    print(f"   Effective batch size: {effective_batch_size}")
-    print(f"   Steps per epoch: {steps_per_epoch}")
-    print(f"   Total steps: {total_steps}")
-    print(f"   Estimated time: ~{total_steps/2.4/60:.1f} minutes")
-    print(f"   Eval every: {eval_steps} steps")
-    
-    args = TrainingArguments(
-        output_dir="./cuad_lora_out",
-        learning_rate=3e-5,  # Slightly lower LR
-        num_train_epochs=num_epochs,
-        per_device_train_batch_size=batch_size,
-        per_device_eval_batch_size=8,
-        gradient_accumulation_steps=gradient_accumulation_steps,
-        fp16=False, bf16=True,
-        eval_strategy="steps",
-        eval_steps=eval_steps,
-        save_steps=save_steps,
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        evaluation_strategy="steps",
+        eval_steps=500,
+        save_steps=500,
         save_total_limit=2,
-        weight_decay=0.01,
-        lr_scheduler_type="cosine",
-        warmup_ratio=0.1,
-        load_best_model_at_end=False,
-        logging_steps=10,  # More frequent logging
-        report_to="none",
-        dataloader_num_workers=2,
-        dataloader_pin_memory=True,
-        remove_unused_columns=True,
+        learning_rate=3e-4,  # LoRA usually supports higher LR
+        per_device_train_batch_size=1,
+        per_device_eval_batch_size=1,
+        num_train_epochs=3,
+        weight_decay=0.0,
+        logging_dir=f"{output_dir}/logs",
+        logging_steps=100,
+        load_best_model_at_end=True,
+        metric_for_best_model="eval_f1",
+        greater_is_better=True,
+        fp16=True,
+        push_to_hub=True,
+        hub_model_id="AvocadoMuffin/deberta_finetuned_qa_lora",
+        hub_strategy="checkpoint",
     )
 
     trainer = Trainer(
         model=model,
-        args=args,
-        train_dataset=train_feats,
-        eval_dataset=val_feats,
-        tokenizer=tok,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        tokenizer=tokenizer,
         data_collator=default_data_collator,
-        compute_metrics=None,
+        compute_metrics=lambda p: compute_metrics(p, tokenizer, datasets["validation"], val_dataset),
     )
 
-    print("🚀 Training…")
-    try:
-        trainer.train()
-        print("✅ Training completed successfully!")
-    except Exception as e:
-        print(f"❌ Training failed: {e}")
-        try:
-            trainer.save_model("./cuad_lora_out_partial")
-            tok.save_pretrained("./cuad_lora_out_partial")
-            print("💾 Partial model saved")
-        except:
-            print("❌ Could not save partial model")
-        raise e
-
-    print("✅ Done. Best eval_loss:", trainer.state.best_metric)
-    trainer.save_model("./cuad_lora_out")
-    tok.save_pretrained("./cuad_lora_out")
-
-    # Push to hub
-    if tokn:
-        for attempt in range(3):
-            try:
-                print(f"⬆️ Pushing to Hub (attempt {attempt + 1}/3)...")
-                trainer.push_to_hub(model_repo, private=False)
-                tok.push_to_hub(model_repo, private=False)
-                print("🚀 Pushed to:", f"https://huggingface.co/{model_repo}")
-                break
-            except Exception as e:
-                print(f"⚠️ Hub push failed: {e}")
-                if attempt < 2:
-                    time.sleep(30)
-                else:
-                    print("💾 Model saved locally (push failed)")
+    trainer.train()
+    trainer.push_to_hub()
 
 if __name__ == "__main__":
-    main()
+    main()