Update train.py

train.py

@@ -3,6 +3,7 @@
 """
 CUAD fine-tune with LoRA - Efficient batch processing version.
 Fixes bottlenecks and uses proper batching instead of chunking.
+GUARANTEED FIX for offset_mapping error.
 """
 
 import os, json, random, gc, time
@@ -21,6 +22,9 @@ from huggingface_hub import login
 
 disable_caching()
 
+# Set tokenizers parallelism to avoid warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 # ─────────────────────────────────────────────────────────────── config ──
 
 MAX_LEN = 384
@@ -108,15 +112,14 @@ def compute_metrics(eval_pred):
 
 # ───────────────────────────────────────────────────────────── preprocessing ──
 
-def preprocess_batch_efficient(examples, tokenizer, is_training=True):
+def preprocess_training_batch(examples, tokenizer):
     """
-    Efficient batch preprocessing using HuggingFace's built-in batch processing.
-    This is much faster than processing examples individually.
+    Training preprocessing - NO offset_mapping included
     """
     questions = examples["question"]
     contexts = examples["context"]
     
-    # Batch tokenization - this is the key efficiency gain
+    # Batch tokenization
     tokenized_examples = tokenizer(
         questions,
         contexts,
@@ -124,19 +127,19 @@ def preprocess_batch_efficient(examples, tokenizer, is_training=True):
         max_length=MAX_LEN,
         stride=DOC_STRIDE,
         return_overflowing_tokens=True,
-        return_offsets_mapping=True,
+        return_offsets_mapping=True,  # We need this temporarily for position calculation
         padding="max_length",
     )
     
     # Map back to original examples
     sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
+    offset_mapping = tokenized_examples.pop("offset_mapping")  # Remove it immediately after use
     
     # Initialize output
     start_positions = []
     end_positions = []
     
-    for i, offsets in enumerate(tokenized_examples["offset_mapping"]):
-        input_ids = tokenized_examples["input_ids"][i]
+    for i, offsets in enumerate(offset_mapping):
         cls_index = 0  # CLS token position
         
         # Get the original example for this tokenized chunk
@@ -176,73 +179,68 @@ def preprocess_batch_efficient(examples, tokenizer, is_training=True):
     tokenized_examples["start_positions"] = start_positions
     tokenized_examples["end_positions"] = end_positions
     
+    # NO offset_mapping or example_id for training
+    return tokenized_examples
+
+def preprocess_validation_batch(examples, tokenizer):
+    """
+    Validation preprocessing - INCLUDES offset_mapping and example_id for post-processing
+    """
+    questions = examples["question"]
+    contexts = examples["context"]
+    
+    # Batch tokenization
+    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",
+    )
+    
+    # Map back to original examples
+    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
+    
     # Add example IDs for evaluation
     tokenized_examples["example_id"] = [
         examples["id"][sample_mapping[i]] for i in range(len(tokenized_examples["input_ids"]))
     ]
     
-    # CRITICAL FIX: Remove offset_mapping for training data
-    # Only keep it for evaluation data
-    if is_training:
-        tokenized_examples.pop("offset_mapping", None)
-    
+    # Keep offset_mapping for post-processing
     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.
-    This is much more memory efficient and faster than manual chunking.
     """
     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_batch_efficient(examples, tokenizer, is_training),
+        lambda examples: preprocess_fn(examples, tokenizer),
         batched=True,
         batch_size=BATCH_SIZE,
         remove_columns=dataset.column_names,
         desc=desc,
-        num_proc=1,  # Use 1 process to avoid memory issues in Spaces
+        num_proc=1,
     )
     
     print(f"✅ {desc} completed: {len(processed)} features")
     return processed
 
-# Custom data collator that ensures no unwanted keys are passed to the model
-class QADataCollator:
-    def __init__(self, tokenizer, training=True):
-        self.tokenizer = tokenizer
-        self.training = training
-        # Keys that the model expects
-        self.model_input_names = ["input_ids", "attention_mask", "token_type_ids", "start_positions", "end_positions"]
-    
-    def __call__(self, features):
-        # Use the default collator first
-        batch = default_data_collator(features)
-        
-        # For training, only keep model input keys
-        if self.training:
-            batch = {k: v for k, v in batch.items() if k in self.model_input_names}
-        else:
-            # For evaluation, keep example_id and offset_mapping for post-processing
-            # but don't pass them to the model
-            model_batch = {k: v for k, v in batch.items() if k in self.model_input_names[:3]}  # No positions for eval
-            # Store extra info for post-processing
-            model_batch["example_id"] = batch.get("example_id")
-            if "offset_mapping" in batch:
-                model_batch["offset_mapping"] = batch["offset_mapping"]
-            batch = model_batch
-            
-        return batch
-
 # ───────────────────────────────────────────────────────────────── main ──
 
 def main():
     global val_raw, val_feats, tok
     
     set_seed(SEED)
-    
-    # Set tokenizers parallelism to avoid the warning
-    os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
     # Model name to store on Hub
     model_repo = os.getenv("MODEL_NAME", "AvocadoMuffin/roberta-cuad-qa-v3")
@@ -282,8 +280,8 @@ def main():
     # LoRA with slightly more aggressive settings for smaller dataset
     lora = LoraConfig(
         task_type=TaskType.QUESTION_ANS,
-        r=32, lora_alpha=64, lora_dropout=0.1,  # Increased for better learning with less data
-        target_modules=["query", "value", "key", "dense"],  # More modules for better coverage
+        r=32, lora_alpha=64, lora_dropout=0.1,
+        target_modules=["query", "value", "key", "dense"],
     )
     model = get_peft_model(model, lora)
     model.print_trainable_parameters()
@@ -291,32 +289,29 @@ def main():
     # ── efficient preprocessing ─────────────────────────────────────────
     print("🔄 Starting efficient preprocessing...")
     
-    # Process training data (remove offset_mapping)
+    # Process training data (NO offset_mapping)
     train_feats = preprocess_dataset_streaming(train_raw, tok, "Training", is_training=True)
     
-    # Process validation data (keep offset_mapping for evaluation)
+    # Process validation data (WITH offset_mapping)
     val_feats = preprocess_dataset_streaming(val_raw, tok, "Validation", is_training=False)
 
     print(f"✅ Preprocessing completed!")
     print(f"   Training features: {len(train_feats)}")
     print(f"   Validation features: {len(val_feats)}")
+    print(f"   Training columns: {train_feats.column_names}")
+    print(f"   Validation columns: {val_feats.column_names}")
 
-    # ── training args with fixed eval/save step alignment ──────────────────
-    # Calculate proper steps that align
+    # ── training args ──────────────────
     batch_size = 16
     gradient_accumulation_steps = 2
     effective_batch_size = batch_size * gradient_accumulation_steps
     
-    # Calculate total training steps
     num_epochs = 6 if USE_SUBSET else 4
     steps_per_epoch = len(train_feats) // effective_batch_size
     total_steps = steps_per_epoch * num_epochs
     
-    # Set eval steps first
-    eval_steps = max(50, steps_per_epoch // 4)  # Evaluate 4 times per epoch
-    
-    # Set save steps as a multiple of eval steps
-    save_steps = eval_steps * 2  # Save every 2 evaluations
+    eval_steps = max(50, steps_per_epoch // 4)
+    save_steps = eval_steps * 2
     
     print(f"📊 Training configuration:")
     print(f"   Steps per epoch: {steps_per_epoch}")
@@ -326,7 +321,7 @@ def main():
     
     args = TrainingArguments(
         output_dir="./cuad_lora_out",
-        learning_rate=5e-5,  # Slightly higher for smaller dataset
+        learning_rate=5e-5,
         num_train_epochs=num_epochs,
         per_device_train_batch_size=batch_size,
         per_device_eval_batch_size=16,
@@ -346,25 +341,18 @@ def main():
         report_to="none",
         dataloader_num_workers=2,
         dataloader_pin_memory=True,
-        remove_unused_columns=False,  # Keep example_id for evaluation
+        remove_unused_columns=True,  # Let the trainer handle column removal
     )
 
-    # Use custom data collators
-    train_collator = QADataCollator(tok, training=True)
-    eval_collator = QADataCollator(tok, training=False)
-
     trainer = Trainer(
         model=model,
         args=args,
         train_dataset=train_feats,
         eval_dataset=val_feats,
         tokenizer=tok,
-        data_collator=train_collator,  # Use custom collator
+        data_collator=default_data_collator,
         compute_metrics=compute_metrics,
     )
-    
-    # Override the evaluation collator
-    trainer.data_collator = eval_collator
 
     print("🚀 Training…")
     try: