From 87edb9c7b9ea35376176784ff735f984077632c0 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sun, 8 Feb 2026 15:27:39 -0800 Subject: [PATCH 1/2] fix crashing --- modeling/dataloader.py | 58 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) diff --git a/modeling/dataloader.py b/modeling/dataloader.py index 5ddf034..a94afe0 100644 --- a/modeling/dataloader.py +++ b/modeling/dataloader.py @@ -47,6 +47,35 @@ def contra_data_collator(mlm_collator, features): code_batch = mlm_collator(code_features) aug_batch = mlm_collator(aug_features) + # Pad to same seq_len (each batch is independently padded to its own max) + code_seq_len = code_batch["input_ids"].size(1) + aug_seq_len = aug_batch["input_ids"].size(1) + if code_seq_len != aug_seq_len: + pad_token_id = mlm_collator.tokenizer.pad_token_id + target_len = max(code_seq_len, aug_seq_len) + if code_seq_len < target_len: + pad = target_len - code_seq_len + code_batch["input_ids"] = torch.nn.functional.pad( + code_batch["input_ids"], (0, pad), value=pad_token_id + ) + code_batch["attention_mask"] = torch.nn.functional.pad( + code_batch["attention_mask"], (0, pad), value=0 + ) + code_batch["labels"] = torch.nn.functional.pad( + code_batch["labels"], (0, pad), value=-100 + ) + else: + pad = target_len - aug_seq_len + aug_batch["input_ids"] = torch.nn.functional.pad( + aug_batch["input_ids"], (0, pad), value=pad_token_id + ) + aug_batch["attention_mask"] = torch.nn.functional.pad( + aug_batch["attention_mask"], (0, pad), value=0 + ) + aug_batch["labels"] = torch.nn.functional.pad( + aug_batch["labels"], (0, pad), value=-100 + ) + # Combine batches batch = { "code_input_ids": code_batch["input_ids"], @@ -138,6 +167,35 @@ def grouped_contra_data_collator(mlm_collator, max_num_augs, features): aug_batch = mlm_collator(aug_features_flat) + # Pad to same seq_len (each batch is independently padded to its own max) + code_seq_len = code_batch["input_ids"].size(1) + aug_seq_len = aug_batch["input_ids"].size(1) + if code_seq_len != aug_seq_len: + pad_token_id = mlm_collator.tokenizer.pad_token_id + target_len = max(code_seq_len, aug_seq_len) + if code_seq_len < target_len: + pad = target_len - code_seq_len + code_batch["input_ids"] = torch.nn.functional.pad( + code_batch["input_ids"], (0, pad), value=pad_token_id + ) + code_batch["attention_mask"] = torch.nn.functional.pad( + code_batch["attention_mask"], (0, pad), value=0 + ) + code_batch["labels"] = torch.nn.functional.pad( + code_batch["labels"], (0, pad), value=-100 + ) + else: + pad = target_len - aug_seq_len + aug_batch["input_ids"] = torch.nn.functional.pad( + aug_batch["input_ids"], (0, pad), value=pad_token_id + ) + aug_batch["attention_mask"] = torch.nn.functional.pad( + aug_batch["attention_mask"], (0, pad), value=0 + ) + aug_batch["labels"] = torch.nn.functional.pad( + aug_batch["labels"], (0, pad), value=-100 + ) + batch = { "code_input_ids": code_batch["input_ids"], "code_attention_mask": code_batch["attention_mask"], From 0bfd833c1b332ca45dc0ee86b80d46a7be7defd0 Mon Sep 17 00:00:00 2001 From: Yifeng He Date: Sun, 8 Feb 2026 19:33:47 -0800 Subject: [PATCH 2/2] fix: disable safetensors saving to handle RoBERTa tied weights --- modeling/pretrain.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modeling/pretrain.py b/modeling/pretrain.py index aad05ef..278635e 100644 --- a/modeling/pretrain.py +++ b/modeling/pretrain.py @@ -330,6 +330,7 @@ def main( save_total_limit=3, load_best_model_at_end=True, dataloader_num_workers=max(1, (os.cpu_count() or 1) // _get_world_size()), + save_safetensors=False, # RoBERTa has tied weights (embeddings ↔ lm_head); safetensors rejects shared tensors ) trainer = ContrastiveTrainer(