diff --git a/docs/ablation.md b/docs/ablation.md new file mode 100644 index 0000000..32a43c3 --- /dev/null +++ b/docs/ablation.md @@ -0,0 +1,111 @@ +# Ablation Studies + +## Goal + +Justify the core design decisions of InvPT for the COLM submission. Each +ablation removes or substitutes exactly one component of the full method to +isolate its contribution. All ablations use **CodeBERT** +(`microsoft/codebert-base`) as the representative backbone to keep compute +manageable. + +## Full Method (Control) + +The full training loss is: + +$$\mathcal{L} = \mathcal{L}_{\text{MLM}}(X) + \mathcal{L}_{\text{MLM}}(X_{\text{inv}}) + \alpha \cdot \mathcal{L}_{\text{SupCon}}(X, X_{\text{inv}})$$ + +with `alpha=1.0`, `temperature=0.1`, `self_contrast=true`. + +Config: `experiments/supcon/codebert.yaml` (run name `InvCodeBERT-supcon`). + +## Ablation Matrix + +| # | Ablation | Question answered | Config | +| --- | ---------------- | --------------------------------------------------------- | -------------------------------- | +| 1 | MLM-only | Does the contrastive loss help at all? | `ablation/mlm_only.yaml` | +| 2 | No self-contrast | Is self-contrast (same code, different masks) beneficial? | `ablation/no_self_contrast.yaml` | +| 3 | + NL (bimodal) | Is PL-only training better than bimodal NL+PL training? | `ablation/include_nl.yaml` | + +### What changes per ablation + +| # | `alpha` | `self_contrast` | `include_nl` | run_name | +| --- | :-----: | :-------------: | :----------: | ---------------------------------- | +| 1 | **0** | true | false | `InvCodeBERT-ablation-mlm-only` | +| 2 | 1.0 | **false** | false | `InvCodeBERT-ablation-no-selfcon` | +| 3 | 1.0 | true | **true** | `InvCodeBERT-ablation-include-nl` | + +## Ablation Details + +### 1 — MLM-Only (`alpha=0`) + +$$\mathcal{L} = \mathcal{L}_{\text{MLM}}(X) + \mathcal{L}_{\text{MLM}}(X_{\text{inv}})$$ + +Removes the contrastive objective entirely. The model still trains on both the +original code $X$ and its invariant augmentation $X_{\text{inv}}$ via MLM, so it +sees the same data as the full method — only the explicit alignment signal is +missing. This isolates the contribution of contrastive learning. + +**Expected outcome.** If the full method outperforms this baseline, the +contrastive loss provides value beyond what MLM on augmented data alone achieves. + +### 2 — No Self-Contrast (`self_contrast=false`) + +Same loss as the full method, but rows without a real code transformation are +**dropped** instead of using the original code as its own augmentation. This +reduces the training set size (rows that failed all transformation operators are +excluded) and removes the "easy" contrastive pairs where both views are the same +code with different MLM masks. + +**Expected outcome.** If the full method outperforms this ablation, self-contrast +is beneficial — it acts as an implicit curriculum (easy pairs stabilize training) +and improves data efficiency. See `doc/self-contrast.md` for a detailed analysis. + +### 3 — Include NL (`include_nl=true`) + +Re-introduces natural language docstrings into the pre-training input, restoring +the bimodal NL+PL setup used by CodeBERT, GraphCodeBERT, and ContraBERT. When +`include_nl=true`, each code input is formatted as `[CLS] [SEP] + [EOS]` instead of the default `[CLS] [EOS]`. The same +prepending applies to the invariant-transformed code. + +This ablation directly tests InvPT's core hypothesis: that PL-only pre-training +is sufficient (and preferable) for learning robust code representations. Prior +work universally relies on NL-PL paired training, and the ICSE'26 submission +showed that adding NL back actually *degraded* performance while consuming more +memory — likely due to overfitting on NL descriptions rather than learning +program semantics. + +**Expected outcome.** The PL-only control should outperform this NL-inclusive +variant, especially on robustness, confirming that docstrings are not needed and +can even be harmful for invariant pre-training. + +**Code change required.** Add an `include_nl: bool = False` parameter to +`PretrainConfig`. In `tokenize_grouped`, when `include_nl=true`, prepend the +docstring to the code before tokenization (for both anchor and augmented +inputs). The docstring field is already carried through the data pipeline but +currently unused during tokenization. + +| File | Change | +| ---------------------- | ------------------------------------------------------------------- | +| `modeling/config.py` | Add `include_nl: bool = False` to `PretrainConfig` | +| `modeling/pretrain.py` | `tokenize_grouped` prepends docstring to code when `include_nl` set | + +## Running + +Smoke-test each config with a 1% sample: + +```bash +# Control +python -m modeling run experiments/supcon/codebert.yaml --sample-rate 0.01 + +# Ablations +python -m modeling run experiments/ablation/mlm_only.yaml --sample-rate 0.01 +python -m modeling run experiments/ablation/no_self_contrast.yaml --sample-rate 0.01 +python -m modeling run experiments/ablation/include_nl.yaml --sample-rate 0.01 +``` + +### What to verify + +- **mlm_only:** Contrastive loss contributes 0 to total loss (alpha=0). +- **no_self_contrast:** Dataset is smaller (filtered rows without augmentation). +- **include_nl:** Tokenized inputs start with docstring before code (`[CLS] docstring [SEP] code [EOS]`). diff --git a/docs/grouped_multi_key_contrast.md b/docs/grouped_multi_key_contrast.md deleted file mode 100644 index a54bff1..0000000 --- a/docs/grouped_multi_key_contrast.md +++ /dev/null @@ -1,155 +0,0 @@ -# Grouped Multi-Key Contrastive Loss - -## Motivation - -In the standard InfoNCE setup (`--contra_mode info_nce`), each training example -is a pair (anchor code, one augmentation). The contrastive loss treats the -diagonal entries of the $B \times B$ similarity matrix as positives and everything -else as negatives. When the dataset contains multiple augmentations of the same -function (e.g. VarRe, RevIf, AA2EA applied to the same source), those -augmentations land in different batch rows and are **pushed apart as negatives** -even though they are semantically equivalent. - -SupCon (`--contra_mode supcon`) mitigates this by masking: it checks -`function_id` at loss time and treats all same-ID embeddings as positives. -However, it relies on same-function augmentations happening to co-occur in the -same mini-batch, which becomes unlikely at small batch sizes. - -Grouped Multi-Key Contrast (`--contra_mode grouped`) solves this structurally: -it **regroups the dataset** so that every batch item already bundles an anchor -with _all_ of its augmentations. Every augmentation is guaranteed to be present -as a positive, regardless of batch size. - -## Algorithm - -### 1. Dataset Regrouping - -Before training, the flat JSONL dataset (one row per `(code, transformed)` pair) -is regrouped by `function_id = SHA-256(code)[:8]`: - -``` -Input (flat): - row 0: code="def foo()...", transformed="def renamed_foo()...", aug_type=VarRe - row 1: code="def foo()...", transformed="def foo_rev()...", aug_type=RevIf - row 2: code="def bar()...", transformed="def bar_v1()...", aug_type=VarRe - -Output (grouped): - group 0: code="def foo()...", transformed_list=["def renamed_foo()...", "def foo_rev()..."] - group 1: code="def bar()...", transformed_list=["def bar_v1()..."] -``` - -Each group's augmentation list is truncated to `max_num_augs` (default 6) and -padded with empty strings to a fixed length for uniform Arrow storage. - -### 2. Collation - -The grouped collator produces: - -| Tensor | Shape | Description | -| ---------------- | ----------------------- | ----------------------------------------------------- | -| `code_input_ids` | $[B, L]$ | Tokenized anchors (MLM-masked) | -| `aug_input_ids` | $[B \cdot K_{\max}, L]$ | Flattened tokenized augmentations (MLM-masked) | -| `group_sizes` | $[B]$ | Number of real (non-padding) augmentations per anchor | - -where $B$ is the batch size, $L$ is `max_seq_length`, and -$K_{\max} = \min(\max_i K_i,\; \texttt{max\_num\_augs})$. - -Padding augmentation slots use `attention_mask = 0` and -`special_tokens_mask = 1` so the MLM collator assigns `labels = -100` to every -token (i.e. padding augmentations contribute zero to the MLM loss). - -### 3. Forward Pass - -A single shared encoder $f_\theta$ processes both anchors and augmentations: - -$$\mathbf{h}_i = f_\theta(\text{code}_i) \quad \text{for } i = 1, \dots, B \qquad \to [B, D]$$ - -$$\mathbf{h}_{i,k} = f_\theta(\text{aug}_{i,k}) \quad \text{for } i = 1, \dots, B,\; k = 1, \dots, K_{\max} \qquad \to [B \cdot K_{\max}, D]$$ - -The CLS token embedding (position 0 of the last hidden layer) is used as the -sequence representation. Both forward passes also produce MLM logits for the -masked language modeling objective. - -### 4. Contrastive Loss - -**Definitions.** Given: - -- Anchor embeddings: $\mathbf{a}_i = \text{normalize}(\mathbf{h}_i)$ for $i = 1, \dots, B$ -- Augmentation embeddings: $\mathbf{v}_{i,k} = \text{normalize}(\mathbf{h}_{i,k})$ for $i = 1, \dots, B$, $k = 1, \dots, K_{\max}$ -- Group sizes: $K_i$ (number of real augmentations for anchor $i$) -- Temperature: $\tau$ -- Similarity function: $\text{sim}(\mathbf{u}, \mathbf{v}) = \mathbf{u}^\top \mathbf{v}$ (dot product of $\ell_2$-normalized vectors) - -**Candidate pool.** For each anchor $i$, the candidate pool consists of: - -- All other anchors $\mathbf{a}_j$ where $j \neq i$ -- All valid (non-padding) augmentation embeddings $\mathbf{v}_{j,k}$ for all $j$ and $k < K_j$ - -**Positive set.** For anchor $i$: $\mathcal{P}(i) = \{\mathbf{v}_{i,k} : k < K_i\}$. - -**Per-positive loss.** For each anchor $i$ and each of its valid positives $\mathbf{v}_{i,k}$: - -$$\ell_{i,k} = -\log \frac{\exp\!\bigl(\text{sim}(\mathbf{a}_i, \mathbf{v}_{i,k}) / \tau\bigr)}{\displaystyle\sum_{\substack{j=1 \\ j \neq i}}^{B} \exp\!\bigl(\text{sim}(\mathbf{a}_i, \mathbf{a}_j) / \tau\bigr) + \sum_{j=1}^{B} \sum_{\substack{k'=0 \\ k' < K_j}}^{K_{\max}-1} \exp\!\bigl(\text{sim}(\mathbf{a}_i, \mathbf{v}_{j,k'}) / \tau\bigr)}$$ - -**Per-anchor loss** (SupCon-style averaging over positives): - -$$\mathcal{L}_i = \frac{1}{K_i} \sum_{k=0}^{K_i - 1} \ell_{i,k}$$ - -**Batch loss** (averaged over anchors with at least one augmentation): - -$$\mathcal{L}_{\text{grouped}} = \frac{1}{|\{i : K_i > 0\}|} \sum_{\substack{i=1 \\ K_i > 0}}^{B} \mathcal{L}_i$$ - -**Note:** The denominator includes the anchor's own augmentations as well. This -differs from some formulations that exclude positives from the denominator; our -version follows the SupCon convention (Khosla et al., 2020) where the -denominator sums over _all_ non-self entries. - -### 5. Total Training Loss - -The total loss combines MLM and contrastive objectives: - -$$\mathcal{L} = \mathcal{L}_{\text{MLM}} + \alpha \cdot \mathcal{L}_{\text{grouped}}$$ - -where the MLM loss averages over anchor and augmentation views: - -$$\mathcal{L}_{\text{MLM}} = \frac{\mathcal{L}_{\text{MLM}}(\text{code}) + \mathcal{L}_{\text{MLM}}(\text{aug})}{2}$$ - -The MLM loss for augmentations is automatically averaged only over non-padding -tokens (padding augmentations have all labels set to $-100$ and contribute zero). - -### 6. Numerical Stability - -The implementation uses the log-sum-exp trick for numerical stability. -Given the raw logits $z_{i,n}$ for anchor $i$ against candidate $n$: - -1. Compute $m_i = \max_{n \notin \text{excluded}} z_{i,n}$, clamped to $\geq 0$ -2. Subtract before exponentiation: $\exp(z_{i,n} - m_i)$ -3. This prevents overflow in $\exp(\cdot)$ when similarity values are large - -## Parameters - -- `--contra_mode` - - default: `info_nce` - - Selects the contrastive loss mode. Set to `grouped` to enable grouped multi-key contrast. -- `--max_num_augs` - - default: `6` - - Maximum augmentations per anchor group ($K_{\max}$). Higher values use more GPU memory per batch item. Set to the number of transformation operators applied to each language (e.g. 3 for Python-only, 6 for Java-only, 6 for mixed). Only used when `--contra_mode grouped`. -- `--alpha` - - default: `1.0` - - Weight $\alpha$ of the contrastive loss relative to MLM. -- `--temperature` - - default: `0.07` - - Contrastive temperature $\tau$. Lower values sharpen the distribution and increase the penalty for hard negatives. -- `--batch_size` - - default: `256` - - Total batch size $B$. Each item now requires $K_{\max}$ augmentation forward passes, so reduce batch size and increase `--gradient_accumulation_steps` compared to `info_nce`/`supcon` modes. - -## Comparison with Other Modes - -| Property | `info_nce` | `supcon` | `grouped` | -| --------------------------------------------- | ---------------------- | -------------------------------------- | ----------------------------------------------------- | -| Positives per anchor | 1 (diagonal) | Variable (depends on batch collisions) | All $K_i$ augmentations (guaranteed) | -| Requires same-function co-occurrence in batch | N/A | Yes | No (grouped at dataset level) | -| Dataset format | Flat (code, aug) pairs | Flat (code, aug) pairs | Grouped (code, $[\text{aug}_1, \dots, \text{aug}_K]$) | -| Memory per batch item | 2 forward passes | 2 forward passes | $1 + K_{\max}$ forward passes | -| Handles variable aug counts | N/A | Naturally (mask-based) | Yes (padding + `group_sizes` mask) | diff --git a/docs/modernbert.md b/docs/modernbert.md index f34cd6f..9da3b6e 100644 --- a/docs/modernbert.md +++ b/docs/modernbert.md @@ -69,7 +69,7 @@ Two new fields control the behavior: - **`pooling`**: Either `"cls"` (default, CLS token) or `"mean"` (mean over non-padding tokens). Mean pooling is recommended for ModernBERT (see below). -All other config fields (`alpha`, `temperature`, `contra_mode`, etc.) work identically. +All other config fields (`alpha`, `temperature`, etc.) work identically. ### CLI Usage @@ -107,18 +107,6 @@ L = L_MLM(code) + L_MLM(aug) + alpha * L_contrastive(code, aug) encoder's last hidden states. Controlled by `alpha` (default 1.0) and `temperature` (default 0.07). -### Contrastive Modes - -All three contrastive modes work with ModernBERT: - -| Mode | Config value | Description | -| ------- | ------------ | --------------------------------------------------------------------------------------------- | -| InfoNCE | `info_nce` | Diagonal positives — each code paired with its single augmentation | -| SupCon | `supcon` | Multi-positive by `function_id` — all augmentations of the same function are mutual positives | -| Grouped | `grouped` | Explicit grouped multi-key contrast with up to `max_num_augs` augmentations per anchor | - -Set via the `contra_mode` field in the YAML config. - ### Self-Contrast Self-contrast (`self_contrast: true`, the default) provides the "easy" curriculum signal. @@ -135,7 +123,7 @@ This is independent of model type and works identically for ModernBERT. When | Aspect | RoBERTa | ModernBERT | | ----------------------- | ----------------------------------------- | ----------------------------- | | Loss function | `L_MLM + alpha * L_contrastive` | Same | -| Contrastive modes | info_nce / supcon / grouped | Same | +| Contrastive loss | SupCon | Same | | Self-contrast | Supported | Same | | MLM masking | 15% via `DataCollatorForLanguageModeling` | Same | | Pooling for contrastive | CLS (default) | Mean (recommended) | diff --git a/experiments/ablation/include_nl.yaml b/experiments/ablation/include_nl.yaml new file mode 100644 index 0000000..c16b56e --- /dev/null +++ b/experiments/ablation/include_nl.yaml @@ -0,0 +1,20 @@ +# Ablation 4: Include NL (bimodal NL+PL training) +# Tests whether PL-only training is better than bimodal NL+PL training + +dataset_path: "data/csn.jsonl" +model_name: "microsoft/codebert-base" + +batch_size: 256 +num_epochs: 3 +gradient_accumulation_steps: 1 +learning_rate: 2.0e-5 + +seed: 0 +run_name: "InvCodeBERT-ablation-include-nl" + +alpha: 1.0 +temperature: 0.1 +max_seq_length: 512 + +self_contrast: true +include_nl: true diff --git a/experiments/ablation/mlm_only.yaml b/experiments/ablation/mlm_only.yaml new file mode 100644 index 0000000..26aaeb0 --- /dev/null +++ b/experiments/ablation/mlm_only.yaml @@ -0,0 +1,19 @@ +# Ablation 1a: MLM-only (no contrastive loss) +# L = L_MLM(X) + L_MLM(X_inv), alpha=0 + +dataset_path: "data/csn.jsonl" +model_name: "microsoft/codebert-base" + +batch_size: 256 +num_epochs: 3 +gradient_accumulation_steps: 1 +learning_rate: 2.0e-5 + +seed: 0 +run_name: "InvCodeBERT-ablation-mlm-only" + +alpha: 0 +temperature: 0.1 +max_seq_length: 512 + +self_contrast: true diff --git a/experiments/ablation/no_self_contrast.yaml b/experiments/ablation/no_self_contrast.yaml new file mode 100644 index 0000000..b1f7b51 --- /dev/null +++ b/experiments/ablation/no_self_contrast.yaml @@ -0,0 +1,19 @@ +# Ablation 2: No self-contrast +# Rows without augmentation are dropped instead of using self-contrast + +dataset_path: "data/csn.jsonl" +model_name: "microsoft/codebert-base" + +batch_size: 256 +num_epochs: 3 +gradient_accumulation_steps: 1 +learning_rate: 2.0e-5 + +seed: 0 +run_name: "InvCodeBERT-ablation-no-selfcon" + +alpha: 1.0 +temperature: 0.1 +max_seq_length: 512 + +self_contrast: false diff --git a/experiments/grouped/codebert.yaml b/experiments/grouped/codebert.yaml deleted file mode 100644 index a1ea518..0000000 --- a/experiments/grouped/codebert.yaml +++ /dev/null @@ -1,22 +0,0 @@ -# InvPT pre-training: Grouped — CodeBERT -# Effective batch size: 16 * 16 = 256 -# Usage: python modeling/cli.py run experiments/grouped/codebert.yaml - -dataset_path: "data/csn.jsonl" -model_name: "microsoft/codebert-base" - -batch_size: 32 -num_epochs: 3 -gradient_accumulation_steps: 8 -learning_rate: 2.0e-5 - -seed: 0 -run_name: "InvCodeBERT-grouped" - -alpha: 1.0 -temperature: 0.1 -max_seq_length: 512 - -contra_mode: "grouped" -max_num_augs: 6 -self_contrast: true diff --git a/experiments/grouped/contrabert_c.yaml b/experiments/grouped/contrabert_c.yaml deleted file mode 100644 index 7d4b740..0000000 --- a/experiments/grouped/contrabert_c.yaml +++ /dev/null @@ -1,23 +0,0 @@ -# InvPT pre-training: Grouped — ContraBERT_C -# Effective batch size: 16 * 16 = 256 -# Usage: python modeling/cli.py run experiments/grouped/contrabert_c.yaml - -dataset_path: "data/csn.jsonl" -model_name: "./saved_models/ContraBERT_C" -tokenizer_name: "microsoft/codebert-base" - -batch_size: 32 -num_epochs: 3 -gradient_accumulation_steps: 8 -learning_rate: 2.0e-5 - -seed: 0 -run_name: "InvContraBERT_C-grouped" - -alpha: 1.0 -temperature: 0.1 -max_seq_length: 512 - -contra_mode: "grouped" -max_num_augs: 6 -self_contrast: true diff --git a/experiments/grouped/contrabert_g.yaml b/experiments/grouped/contrabert_g.yaml deleted file mode 100644 index 1c7cfd1..0000000 --- a/experiments/grouped/contrabert_g.yaml +++ /dev/null @@ -1,23 +0,0 @@ -# InvPT pre-training: Grouped — ContraBERT_G -# Effective batch size: 16 * 16 = 256 -# Usage: python modeling/cli.py run experiments/grouped/contrabert_g.yaml - -dataset_path: "data/csn.jsonl" -model_name: "./saved_models/ContraBERT_G" -tokenizer_name: "microsoft/graphcodebert-base" - -batch_size: 32 -num_epochs: 3 -gradient_accumulation_steps: 8 -learning_rate: 2.0e-5 - -seed: 0 -run_name: "InvContraBERT_G-grouped" - -alpha: 1.0 -temperature: 0.1 -max_seq_length: 512 - -contra_mode: "grouped" -max_num_augs: 6 -self_contrast: true diff --git a/experiments/grouped/graphcodebert.yaml b/experiments/grouped/graphcodebert.yaml deleted file mode 100644 index d7eb62a..0000000 --- a/experiments/grouped/graphcodebert.yaml +++ /dev/null @@ -1,22 +0,0 @@ -# InvPT pre-training: Grouped — GraphCodeBERT -# Effective batch size: 16 * 16 = 256 -# Usage: python modeling/cli.py run experiments/grouped/graphcodebert.yaml - -dataset_path: "data/csn.jsonl" -model_name: "microsoft/graphcodebert-base" - -batch_size: 32 -num_epochs: 3 -gradient_accumulation_steps: 8 -learning_rate: 2.0e-5 - -seed: 0 -run_name: "InvGraphCodeBERT-grouped" - -alpha: 1.0 -temperature: 0.1 -max_seq_length: 512 - -contra_mode: "grouped" -max_num_augs: 6 -self_contrast: true diff --git a/experiments/grouped/modernbert.yaml b/experiments/grouped/modernbert.yaml deleted file mode 100644 index 93a1251..0000000 --- a/experiments/grouped/modernbert.yaml +++ /dev/null @@ -1,24 +0,0 @@ -# InvPT pre-training: Grouped — ModernBERT-base -# Effective batch size: 16 * 16 = 256 -# Usage: python modeling/cli.py run experiments/grouped/modernbert.yaml - -dataset_path: "data/csn.jsonl" -model_name: "answerdotai/ModernBERT-base" -model_type: "modernbert" -pooling: "mean" - -batch_size: 32 -num_epochs: 3 -gradient_accumulation_steps: 8 -learning_rate: 2.0e-5 - -seed: 0 -run_name: "InvModernBERT-grouped" - -alpha: 1.0 -temperature: 0.1 -max_seq_length: 512 - -contra_mode: "grouped" -max_num_augs: 6 -self_contrast: true diff --git a/experiments/supcon/codebert.yaml b/experiments/supcon/codebert.yaml index 8dca2c0..4cbbaa8 100644 --- a/experiments/supcon/codebert.yaml +++ b/experiments/supcon/codebert.yaml @@ -1,14 +1,14 @@ -# InvPT pre-training: SupCon — CodeBERT -# Effective batch size: 64 * 4 = 256 +# InvPT pre-training: SupCon — CodeBERT (pilot: larger batch + longer training) +# Effective batch size: 1024 (4x baseline) # Usage: python modeling/cli.py run experiments/supcon/codebert.yaml dataset_path: "data/csn.jsonl" model_name: "microsoft/codebert-base" -batch_size: 256 -num_epochs: 3 +batch_size: 1024 +num_epochs: 10 gradient_accumulation_steps: 1 -learning_rate: 2.0e-5 +learning_rate: 5.0e-6 seed: 0 run_name: "InvCodeBERT-supcon" @@ -17,5 +17,4 @@ alpha: 1.0 temperature: 0.1 max_seq_length: 512 -contra_mode: "supcon" self_contrast: true diff --git a/experiments/supcon/contrabert_c.yaml b/experiments/supcon/contrabert_c.yaml index 14d4c56..9ff2ca0 100644 --- a/experiments/supcon/contrabert_c.yaml +++ b/experiments/supcon/contrabert_c.yaml @@ -18,5 +18,4 @@ alpha: 1.0 temperature: 0.1 max_seq_length: 512 -contra_mode: "supcon" self_contrast: true diff --git a/experiments/supcon/contrabert_g.yaml b/experiments/supcon/contrabert_g.yaml index f3fe39a..375d2b9 100644 --- a/experiments/supcon/contrabert_g.yaml +++ b/experiments/supcon/contrabert_g.yaml @@ -18,5 +18,4 @@ alpha: 1.0 temperature: 0.1 max_seq_length: 512 -contra_mode: "supcon" self_contrast: true diff --git a/experiments/supcon/graphcodebert.yaml b/experiments/supcon/graphcodebert.yaml index e5fc701..85801cb 100644 --- a/experiments/supcon/graphcodebert.yaml +++ b/experiments/supcon/graphcodebert.yaml @@ -17,5 +17,4 @@ alpha: 1.0 temperature: 0.1 max_seq_length: 512 -contra_mode: "supcon" self_contrast: true diff --git a/experiments/supcon/modernbert.yaml b/experiments/supcon/modernbert.yaml index 0f23060..051b583 100644 --- a/experiments/supcon/modernbert.yaml +++ b/experiments/supcon/modernbert.yaml @@ -2,7 +2,7 @@ # Effective batch size: 64 * 4 = 256 # Usage: python modeling/cli.py run experiments/supcon/modernbert.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "answerdotai/ModernBERT-base" model_type: "modernbert" pooling: "mean" @@ -19,5 +19,4 @@ alpha: 1.0 temperature: 0.1 max_seq_length: 512 -contra_mode: "supcon" self_contrast: true diff --git a/experiments_downstream/run_all_downstream.py b/experiments_downstream/run_all_downstream.py index aa0694e..c140460 100644 --- a/experiments_downstream/run_all_downstream.py +++ b/experiments_downstream/run_all_downstream.py @@ -80,6 +80,35 @@ class RunHandle: ), } +# Ablation models — all based on CodeBERT +ABLATION_MODELS: dict[str, ModelSpec] = { + "contra-only": ModelSpec( + "./saved_models/InvCodeBERT-ablation-contra-only/final", + "microsoft/codebert-base", + "roberta", + ), + "mlm-only": ModelSpec( + "./saved_models/InvCodeBERT-ablation-mlm-only/final", + "microsoft/codebert-base", + "roberta", + ), + "no-self-contrast": ModelSpec( + "./saved_models/InvCodeBERT-ablation-no-selfcon/final", + "microsoft/codebert-base", + "roberta", + ), + "infonce": ModelSpec( + "./saved_models/InvCodeBERT-ablation-infonce/final", + "microsoft/codebert-base", + "roberta", + ), + "include-nl": ModelSpec( + "./saved_models/InvCodeBERT-ablation-include-nl/final", + "microsoft/codebert-base", + "roberta", + ), +} + TASKS = [ ("Clone-detection-POJ104", None), @@ -193,69 +222,16 @@ def run_task( return RunHandle(label=label, process=proc, log_file=log_file) -@app.command() -def main( - all_models: bool = typer.Option( - False, "--all", help="Run all models in the registry." - ), - loss: str = typer.Option( - None, "--loss", help="Training loss identifier (e.g., supcon)." - ), - model: str = typer.Option( - None, - "--model", - help="Pretrained model key(s), comma-separated (e.g., inv-codebert,inv-graphcodebert).", - ), - gpus: str = typer.Option( - "0,1,2,3,4,5,6,7", - "--gpus", - help="Comma-separated GPU ids to use.", - ), - results_root: str = typer.Option( - "results", "--results-root", help="Base output directory for results." - ), - dry_run: bool = typer.Option( - False, "--dry-run", help="Print commands without executing them." - ), +def _run_jobs( + jobs: list[tuple[str, ModelSpec, str, str | None]], + gpu_ids: list[str], + results_root_path: Path, + dry_run: bool, + desc: str = "Downstream", ) -> None: + """Execute *jobs* across *gpu_ids* with a work-stealing thread pool.""" root = Path(__file__).resolve().parents[1] - - # Resolve which models to run - if all_models: - if model is not None: - raise typer.BadParameter("Cannot use --model with --all") - entries = list(MODELS.items()) - if loss is not None: - loss_key = loss.strip() - entries = [((m, lk), s) for (m, lk), s in entries if lk == loss_key] - if not entries: - raise typer.BadParameter(f"No models found for loss={loss}") - else: - if model is None or loss is None: - raise typer.BadParameter( - "Either --all or both --model and --loss are required" - ) - loss_key = loss.strip() - model_keys = [m.strip() for m in model.split(",") if m.strip()] - if not model_keys: - raise typer.BadParameter("No model keys provided") - entries = [] - for model_key in model_keys: - spec = resolve_model(model_key, loss_key) - entries.append(((model_key, loss_key), spec)) - - gpu_ids = [gpu.strip() for gpu in gpus.split(",") if gpu.strip()] - if not gpu_ids: - raise typer.BadParameter("No GPU ids provided") - - # Build all jobs: (model_key, spec, task_dir, subset) - jobs: list[tuple[str, ModelSpec, str, str | None]] = [] - for (mk, _lk), sp in entries: - for task_dir, subset in TASKS: - jobs.append((mk, sp, task_dir, subset)) - total = len(jobs) - results_root_path = Path(results_root).resolve() if dry_run: for i, (mk, sp, task_dir, subset) in enumerate(jobs): @@ -264,7 +240,6 @@ def main( print(f"\n[dry-run] {total} total jobs across {len(gpu_ids)} GPUs") raise typer.Exit(0) - # GPU work-stealing pool gpu_pool: queue.Queue[str] = queue.Queue() for gid in gpu_ids: gpu_pool.put(gid) @@ -273,7 +248,7 @@ def main( failed = 0 lock = threading.Lock() failures: list[str] = [] - pbar = tqdm(total=total, desc="Downstream", unit="task") + pbar = tqdm(total=total, desc=desc, unit="task") def run_job(mk: str, sp: ModelSpec, task_dir: str, subset: str | None) -> None: nonlocal running, failed @@ -317,7 +292,127 @@ def run_job(mk: str, sp: ModelSpec, task_dir: str, subset: str | None) -> None: print(f" - {failure}") raise typer.Exit(1) - print(f"\n[done] All {total} downstream tasks completed successfully.") + print(f"\n[done] All {total} tasks completed successfully.") + + +def _parse_gpus(gpus: str) -> list[str]: + gpu_ids = [gpu.strip() for gpu in gpus.split(",") if gpu.strip()] + if not gpu_ids: + raise typer.BadParameter("No GPU ids provided") + return gpu_ids + + +@app.command() +def run( + all_models: bool = typer.Option( + False, "--all", help="Run all models in the registry." + ), + loss: str = typer.Option( + None, "--loss", help="Training loss identifier (e.g., supcon)." + ), + model: str = typer.Option( + None, + "--model", + help="Pretrained model key(s), comma-separated (e.g., inv-codebert,inv-graphcodebert).", + ), + gpus: str = typer.Option( + "0,1,2,3,4,5,6,7", + "--gpus", + help="Comma-separated GPU ids to use.", + ), + results_root: str = typer.Option( + "results", "--results-root", help="Base output directory for results." + ), + dry_run: bool = typer.Option( + False, "--dry-run", help="Print commands without executing them." + ), +) -> None: + """Run downstream evaluation for pretrained / baseline models.""" + # Resolve which models to run + if all_models: + if model is not None: + raise typer.BadParameter("Cannot use --model with --all") + entries = list(MODELS.items()) + if loss is not None: + loss_key = loss.strip() + entries = [((m, lk), s) for (m, lk), s in entries if lk == loss_key] + if not entries: + raise typer.BadParameter(f"No models found for loss={loss}") + else: + if model is None or loss is None: + raise typer.BadParameter( + "Either --all or both --model and --loss are required" + ) + loss_key = loss.strip() + model_keys = [m.strip() for m in model.split(",") if m.strip()] + if not model_keys: + raise typer.BadParameter("No model keys provided") + entries = [] + for model_key in model_keys: + spec = resolve_model(model_key, loss_key) + entries.append(((model_key, loss_key), spec)) + + gpu_ids = _parse_gpus(gpus) + + jobs: list[tuple[str, ModelSpec, str, str | None]] = [] + for (mk, _lk), sp in entries: + for task_dir, subset in TASKS: + jobs.append((mk, sp, task_dir, subset)) + + _run_jobs(jobs, gpu_ids, Path(results_root).resolve(), dry_run) + + +@app.command() +def ablation( + all_models: bool = typer.Option(False, "--all", help="Run all ablation models."), + model: str = typer.Option( + None, + "--model", + help=( + "Ablation model key(s), comma-separated. " + f"Available: {', '.join(sorted(ABLATION_MODELS))}." + ), + ), + gpus: str = typer.Option( + "0,1,2,3,4,5,6,7", + "--gpus", + help="Comma-separated GPU ids to use.", + ), + results_root: str = typer.Option( + "results", "--results-root", help="Base output directory for results." + ), + dry_run: bool = typer.Option( + False, "--dry-run", help="Print commands without executing them." + ), +) -> None: + """Run downstream evaluation for ablation models.""" + if all_models: + if model is not None: + raise typer.BadParameter("Cannot use --model with --all") + selected = list(ABLATION_MODELS.items()) + else: + if model is None: + raise typer.BadParameter("Either --all or --model is required") + keys = [k.strip() for k in model.split(",") if k.strip()] + if not keys: + raise typer.BadParameter("No model keys provided") + selected = [] + for k in keys: + if k not in ABLATION_MODELS: + available = ", ".join(sorted(ABLATION_MODELS)) + raise typer.BadParameter( + f"Unknown ablation model: {k}. Available: {available}" + ) + selected.append((k, ABLATION_MODELS[k])) + + gpu_ids = _parse_gpus(gpus) + + jobs: list[tuple[str, ModelSpec, str, str | None]] = [] + for mk, sp in selected: + for task_dir, subset in TASKS: + jobs.append((mk, sp, task_dir, subset)) + + _run_jobs(jobs, gpu_ids, Path(results_root).resolve(), dry_run, desc="Ablation") if __name__ == "__main__": diff --git a/modeling/_types.py b/modeling/_types.py index fff7ceb..a47afdb 100644 --- a/modeling/_types.py +++ b/modeling/_types.py @@ -1,12 +1,6 @@ from enum import Enum -class ContraMode(str, Enum): - INFO_NCE = "info_nce" - SUPCON = "supcon" - GROUPED = "grouped" - - class ModelType(str, Enum): ROBERTA = "roberta" MODERNBERT = "modernbert" diff --git a/modeling/cli.py b/modeling/cli.py index 09fad08..46b4915 100644 --- a/modeling/cli.py +++ b/modeling/cli.py @@ -4,7 +4,7 @@ import typer -from modeling._types import ContraMode, ModelType +from modeling._types import ModelType from modeling.common import default_num_proc from modeling.config import load_config from modeling.pretrain import main @@ -47,6 +47,53 @@ def run( main(**kwargs) +@app.command("run-all") +def run_all( + config_dir: Annotated[ + Path, + typer.Argument(help="Directory containing YAML experiment config files."), + ], + sample_rate: Annotated[ + Optional[float], + typer.Option("--sample-rate", "-sr", help="Override sample rate from config."), + ] = None, + self_contrast: Annotated[ + Optional[bool], + typer.Option( + help="Override self-contrast setting from config. Use --no-self-contrast to disable.", + ), + ] = None, +) -> None: + """Run pre-training for every YAML config in a directory, sequentially. + + Example: python -m modeling run-all experiments/supcon/ + """ + configs = sorted(config_dir.glob("*.yaml")) + if not configs: + typer.echo(f"No .yaml files found in {config_dir}") + raise typer.Exit(1) + + typer.echo(f"Found {len(configs)} config(s) in {config_dir}:") + for c in configs: + typer.echo(f" - {c}") + typer.echo("") + + for i, cfg_path in enumerate(configs, 1): + typer.echo(f"{'=' * 42}") + typer.echo(f"[{i}/{len(configs)}] Running: {cfg_path}") + typer.echo(f"{'=' * 42}") + cfg = load_config(cfg_path) + kwargs = asdict(cfg) + if sample_rate is not None: + kwargs["sample_rate"] = sample_rate + if self_contrast is not None: + kwargs["self_contrast"] = self_contrast + main(**kwargs) + typer.echo("") + + typer.echo("All experiments complete.") + + @app.command() def pretrain( dataset_path: Annotated[ @@ -90,12 +137,6 @@ def pretrain( checkpoint: Annotated[ Optional[str], typer.Option(help="Path to checkpoint for weight loading.") ] = None, - contra_mode: Annotated[ - ContraMode, typer.Option(help="Contrastive loss mode.") - ] = ContraMode.INFO_NCE, - max_num_augs: Annotated[ - int, typer.Option(help="Max augmentations per anchor (grouped mode).") - ] = 6, self_contrast: Annotated[ bool, typer.Option( @@ -109,6 +150,11 @@ def pretrain( str, typer.Option(help="Pooling strategy for contrastive embeddings (cls or mean)."), ] = "cls", + mlm_weight: Annotated[float, typer.Option(help="Weight for MLM loss.")] = 1.0, + include_nl: Annotated[ + bool, + typer.Option(help="Include NL docstrings in input (bimodal NL+PL training)."), + ] = False, ) -> None: """Run pre-training with all parameters specified as CLI options. @@ -132,11 +178,11 @@ def pretrain( num_proc=num_proc, resume=resume, checkpoint=checkpoint, - contra_mode=contra_mode, - max_num_augs=max_num_augs, self_contrast=self_contrast, model_type=model_type, pooling=pooling, + mlm_weight=mlm_weight, + include_nl=include_nl, ) diff --git a/modeling/config.py b/modeling/config.py index 7f84ec2..9d93ece 100644 --- a/modeling/config.py +++ b/modeling/config.py @@ -4,10 +4,10 @@ import dacite import yaml -from ._types import ContraMode, ModelType +from ._types import ModelType from .common import default_num_proc -_DACITE_CONFIG = dacite.Config(cast=[ContraMode, ModelType]) +_DACITE_CONFIG = dacite.Config(cast=[ModelType]) @dataclass @@ -25,16 +25,16 @@ class PretrainConfig: learning_rate: float = 2e-4 resume: bool = False alpha: float = 1.0 + mlm_weight: float = 1.0 temperature: float = 0.07 max_seq_length: int = 256 sample_rate: float = 1.0 checkpoint: str | None = None tokenizer_name: str | None = None - contra_mode: ContraMode = ContraMode.INFO_NCE - max_num_augs: int = 6 self_contrast: bool = True model_type: ModelType = ModelType.ROBERTA pooling: str = "cls" + include_nl: bool = False def load_config(path: str | Path) -> PretrainConfig: diff --git a/modeling/dataloader.py b/modeling/dataloader.py index a94afe0..54ce3af 100644 --- a/modeling/dataloader.py +++ b/modeling/dataloader.py @@ -93,122 +93,3 @@ def contra_data_collator(mlm_collator, features): ) return batch - - -def grouped_contra_data_collator(mlm_collator, max_num_augs, features): - """Collate grouped samples where each item has 1 anchor + variable-count augmentations. - - Each feature dict contains: - - code_input_ids, code_attention_mask, code_special_tokens_mask (anchor) - - aug_input_ids_list: list of K token-id lists (variable K per sample) - - aug_attention_mask_list, aug_special_tokens_mask_list: same structure - - function_id: int - - Returns a batch dict with: - - code_input_ids: [B, seq_len] - - code_attention_mask: [B, seq_len] - - code_labels: [B, seq_len] (MLM labels for anchors) - - aug_input_ids: [B * max_K, seq_len] (flattened augmentations) - - aug_attention_mask: [B * max_K, seq_len] - - aug_labels: [B * max_K, seq_len] (MLM labels; padding has -100) - - group_sizes: [B] (real aug count per anchor) - - function_id: [B] - """ - # --- Anchor collation (same as existing) --- - code_features = [ - { - "input_ids": f["code_input_ids"], - "attention_mask": f["code_attention_mask"], - "special_tokens_mask": f["code_special_tokens_mask"], - } - for f in features - ] - code_batch = mlm_collator(code_features) - - # --- Augmentation collation --- - seq_len = len(features[0]["code_input_ids"]) - - # Determine max_K for this batch (capped by max_num_augs) - max_K = min( - max(len(f["aug_input_ids_list"]) for f in features), - max_num_augs, - ) - # Ensure at least 1 slot to avoid empty tensors - max_K = max(max_K, 1) - - group_sizes = [] - aug_features_flat = [] - - for f in features: - aug_ids = f["aug_input_ids_list"][:max_K] - aug_masks = f["aug_attention_mask_list"][:max_K] - aug_special = f["aug_special_tokens_mask_list"][:max_K] - real_K = len(aug_ids) - group_sizes.append(real_K) - - for k in range(real_K): - aug_features_flat.append( - { - "input_ids": aug_ids[k], - "attention_mask": aug_masks[k], - "special_tokens_mask": aug_special[k], - } - ) - - # Pad remaining slots with zeros (special_tokens_mask=1 → no MLM) - for _ in range(max_K - real_K): - aug_features_flat.append( - { - "input_ids": [0] * seq_len, - "attention_mask": [0] * seq_len, - "special_tokens_mask": [1] * seq_len, - } - ) - - 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"], - "code_labels": code_batch["labels"], - "aug_input_ids": aug_batch["input_ids"], # [B * max_K, seq_len] - "aug_attention_mask": aug_batch["attention_mask"], - "aug_labels": aug_batch["labels"], - "group_sizes": torch.tensor(group_sizes, dtype=torch.long), - } - - if "function_id" in features[0]: - batch["function_id"] = torch.tensor( - [f["function_id"] for f in features], dtype=torch.long - ) - - return batch diff --git a/modeling/model.py b/modeling/model.py index 5057451..b7c7ae9 100644 --- a/modeling/model.py +++ b/modeling/model.py @@ -3,8 +3,6 @@ import torch.nn.functional as F from transformers import Trainer -from ._types import ContraMode - class SplitHeadWrapper(nn.Module): """Wraps a ``*ForMaskedLM`` model so the LM head is applied per-chunk. @@ -110,16 +108,6 @@ def forward( return mlm_loss, last_hidden -def info_nce_loss(query, key, temperature=0.07): - device = query.device - query = F.normalize(query, dim=1) - key = F.normalize(key, dim=1) - logits = torch.matmul(query, key.transpose(-1, -2)) / temperature - labels = torch.arange(query.size(0)).long().to(device) - loss = F.cross_entropy(logits, labels) - return loss - - def barlow_twins_loss(query, key, lambda_param=0.005): # Normalize representations along batch dimension query = (query - query.mean(dim=0)) / query.std(dim=0) @@ -219,103 +207,6 @@ def supcon_loss( return loss -def grouped_contrastive_loss( - anchor_embeddings: torch.Tensor, - aug_embeddings: torch.Tensor, - group_sizes: torch.Tensor, - temperature: float = 0.07, - eps: float = 1e-8, -) -> torch.Tensor: - """Grouped multi-key contrastive loss. - - Each anchor has a variable number of augmentation positives (given by - ``group_sizes``). Negatives are all other anchors and their augmentations. - - Uses per-positive log-prob averaging (SupCon-style) with log-sum-exp - stabilization. - - Args: - anchor_embeddings: ``[B, D]`` CLS embeddings of anchor codes. - aug_embeddings: ``[B * max_K, D]`` CLS embeddings of flattened - augmentations (padded groups have zero-vectors). - group_sizes: ``[B]`` number of real augmentations per anchor. - temperature: contrastive temperature. - eps: numerical stability constant. - - Returns: - Scalar loss averaged over anchors that have at least one augmentation. - """ - device = anchor_embeddings.device - B = anchor_embeddings.size(0) - total_augs = aug_embeddings.size(0) - max_K = total_augs // B - - # Normalize - anchor_embeddings = F.normalize(anchor_embeddings, dim=1) # [B, D] - aug_embeddings = F.normalize(aug_embeddings, dim=1) # [B*max_K, D] - - # Reshape aug embeddings to [B, max_K, D] - aug_reshaped = aug_embeddings.view(B, max_K, -1) - - # Validity mask: [B, max_K] — True for real augmentations - arange_K = torch.arange(max_K, device=device).unsqueeze(0) # [1, max_K] - valid_mask = arange_K < group_sizes.unsqueeze(1) # [B, max_K] - - # --- Similarity matrices --- - # anchor-to-anchor: [B, B] - sim_a2a = torch.matmul(anchor_embeddings, anchor_embeddings.T) / temperature - # anchor-to-all-augs: [B, B*max_K] - sim_a2aug = torch.matmul(anchor_embeddings, aug_embeddings.T) / temperature - - # --- Build denominator exclusion mask [B, B + B*max_K] --- - # Exclude: self-anchor (diagonal of a2a) + padding aug positions - all_logits = torch.cat([sim_a2a, sim_a2aug], dim=1) # [B, B + B*max_K] - - # Self-anchor exclusion - self_anchor_mask = torch.eye(B, dtype=torch.bool, device=device) # [B, B] - - # Padding aug exclusion: [B, B*max_K] - aug_valid_global = valid_mask.reshape(-1) # [B*max_K] - aug_invalid_mask = ~aug_valid_global.unsqueeze(0).expand(B, -1) # [B, B*max_K] - - denom_exclude = torch.cat([self_anchor_mask, aug_invalid_mask], dim=1) - - # Log-sum-exp stability - max_logit = ( - all_logits.masked_fill(denom_exclude, float("-inf")) - .max(dim=1, keepdim=True) - .values - ) - max_logit = max_logit.clamp(min=0.0) - - exp_logits = torch.exp(all_logits - max_logit) - exp_logits = exp_logits.masked_fill(denom_exclude, 0.0) - denom = exp_logits.sum(dim=1, keepdim=True) + eps # [B, 1] - - # --- Positive logits: anchor i vs its own augmentations --- - # sim_a2aug reshaped to [B, B, max_K] — index [i, i, :] = anchor i's augs - sim_a2aug_grouped = sim_a2aug.view(B, B, max_K) - pos_logits = sim_a2aug_grouped[ - torch.arange(B, device=device), torch.arange(B, device=device), : - ] # [B, max_K] - - # Per-positive log-prob - exp_pos = torch.exp(pos_logits - max_logit) # [B, max_K] - log_prob_pos = torch.log(exp_pos / denom + eps) # [B, max_K] - log_prob_pos = log_prob_pos.masked_fill(~valid_mask, 0.0) - - # Average over valid positives per anchor, then over anchors - per_anchor_loss = -log_prob_pos.sum(dim=1) / group_sizes.float().clamp(min=1.0) - - has_augs = group_sizes > 0 - if has_augs.any(): - loss = per_anchor_loss[has_augs].mean() - else: - loss = torch.tensor(0.0, device=device, requires_grad=True) - - return loss - - class ContrastiveTrainer(Trainer): """HF Trainer subclass for contrastive pre-training. @@ -328,16 +219,16 @@ class ContrastiveTrainer(Trainer): def __init__( self, alpha=1.0, + mlm_weight=1.0, temperature=0.07, - contra_mode="info_nce", pooling="cls", *args, **kwargs, ): super().__init__(*args, **kwargs) self.alpha = alpha + self.mlm_weight = mlm_weight self.temperature = temperature - self.contra_mode = ContraMode(contra_mode) self.pooling = pooling def _pool( @@ -363,9 +254,6 @@ def _pool( def compute_loss( self, model, inputs, return_outputs=False, num_items_in_batch=None ): - if self.contra_mode == ContraMode.GROUPED: - return self._compute_grouped_loss(model, inputs, return_outputs) - device = model.device code_input_ids = inputs["code_input_ids"].to(device) code_attention_mask = inputs["code_attention_mask"].to(device) @@ -392,66 +280,15 @@ def compute_loss( code_embeddings = self._pool(last_hidden[:B], code_attention_mask) aug_embeddings = self._pool(last_hidden[B:], aug_attention_mask) - # Compute contrastive loss between code and its augmentation - if self.contra_mode == ContraMode.SUPCON: - all_embeddings = torch.cat([code_embeddings, aug_embeddings], dim=0) - function_ids = inputs["function_id"].to(device) - all_function_ids = torch.cat([function_ids, function_ids], dim=0) - contrastive_loss = supcon_loss( - all_embeddings, all_function_ids, self.temperature - ) - else: - contrastive_loss = info_nce_loss( - code_embeddings, - aug_embeddings, - self.temperature, - ) - - total_loss = mlm_loss + self.alpha * contrastive_loss - - return (total_loss, None) if return_outputs else total_loss - - def _compute_grouped_loss(self, model, inputs, return_outputs=False): - """Compute loss for grouped multi-key contrast mode. - - Inputs contain: - - code_input_ids: [B, seq_len] - - code_attention_mask, code_labels: same shape - - aug_input_ids: [B * max_K, seq_len] - - aug_attention_mask, aug_labels: same shape - - group_sizes: [B] - """ - device = model.device - code_input_ids = inputs["code_input_ids"].to(device) - code_attention_mask = inputs["code_attention_mask"].to(device) - code_labels = inputs["code_labels"].to(device) - aug_input_ids = inputs["aug_input_ids"].to(device) - aug_attention_mask = inputs["aug_attention_mask"].to(device) - aug_labels = inputs["aug_labels"].to(device) - group_sizes = inputs["group_sizes"].to(device) - - B = code_input_ids.size(0) - - all_input_ids = torch.cat([code_input_ids, aug_input_ids], dim=0) - all_attention_mask = torch.cat([code_attention_mask, aug_attention_mask], dim=0) - - mlm_loss, last_hidden = model( - input_ids=all_input_ids, - attention_mask=all_attention_mask, - labels_a=code_labels, - labels_b=aug_labels, - split_at=B, - ) - - code_embeddings = self._pool(last_hidden[:B], code_attention_mask) - aug_embeddings = self._pool(last_hidden[B:], aug_attention_mask) - - # Contrastive loss - contrastive_loss = grouped_contrastive_loss( - code_embeddings, aug_embeddings, group_sizes, self.temperature + # Supervised contrastive loss (SupCon) + all_embeddings = torch.cat([code_embeddings, aug_embeddings], dim=0) + function_ids = inputs["function_id"].to(device) + all_function_ids = torch.cat([function_ids, function_ids], dim=0) + contrastive_loss = supcon_loss( + all_embeddings, all_function_ids, self.temperature ) - total_loss = mlm_loss + self.alpha * contrastive_loss + total_loss = self.mlm_weight * mlm_loss + self.alpha * contrastive_loss return (total_loss, None) if return_outputs else total_loss diff --git a/modeling/pretrain.py b/modeling/pretrain.py index 56d17d7..f0f19cb 100644 --- a/modeling/pretrain.py +++ b/modeling/pretrain.py @@ -2,11 +2,10 @@ import hashlib import os -from collections import defaultdict from functools import partial from accelerate import PartialState -from datasets import Dataset, Features, Value, load_dataset +from datasets import Features, Value, load_dataset from transformers import ( AutoConfig, AutoModelForMaskedLM, @@ -15,9 +14,8 @@ TrainingArguments, ) -from ._types import ContraMode from .common import default_num_proc, set_seed -from .dataloader import contra_data_collator, grouped_contra_data_collator +from .dataloader import contra_data_collator from .model import ContrastiveTrainer, SplitHeadWrapper @@ -32,19 +30,36 @@ def compute_function_id(code: str) -> int: return int.from_bytes(digest[:8], "big") & 0x7FFFFFFFFFFFFFFF -def tokenize(tokenizer, example, max_seq_length=256): - code_inputs = tokenizer( - example["code"], - truncation=True, - max_length=max_seq_length, - return_special_tokens_mask=True, - ) - aug_inputs = tokenizer( - example["transformed"], - truncation=True, - max_length=max_seq_length, - return_special_tokens_mask=True, - ) +def tokenize(tokenizer, example, max_seq_length=256, include_nl=False): + if include_nl: + # Bimodal: [CLS] docstring [SEP] code [EOS] + code_inputs = tokenizer( + example["docstring"], + example["code"], + truncation=True, + max_length=max_seq_length, + return_special_tokens_mask=True, + ) + aug_inputs = tokenizer( + example["docstring"], + example["transformed"], + truncation=True, + max_length=max_seq_length, + return_special_tokens_mask=True, + ) + else: + code_inputs = tokenizer( + example["code"], + truncation=True, + max_length=max_seq_length, + return_special_tokens_mask=True, + ) + aug_inputs = tokenizer( + example["transformed"], + truncation=True, + max_length=max_seq_length, + return_special_tokens_mask=True, + ) result = { "code_input_ids": code_inputs["input_ids"], "code_attention_mask": code_inputs["attention_mask"], @@ -62,148 +77,6 @@ def tokenize(tokenizer, example, max_seq_length=256): return result -def regroup_dataset(dataset, max_num_augs: int = 6) -> Dataset: - """Regroup flat (code, transformed) rows by function_id into grouped records. - - Each output record contains one anchor code and a list of its augmentations, - padded to ``max_num_augs`` with empty strings. Groups with zero successful - augmentations are filtered out. - - Returns: - A ``datasets.Dataset`` with columns: repo, func_name, language, code, - docstring, transformed_list, aug_type_list, num_augs, function_id. - """ - groups: dict[int, dict] = defaultdict( - lambda: { - "repo": None, - "func_name": None, - "language": None, - "code": None, - "docstring": None, - "transformed_list": [], - "aug_type_list": [], - } - ) - - # Iterating a HF Dataset is much faster than random indexing (dataset[i]). - for row in dataset: - fid = compute_function_id(row["code"]) - g = groups[fid] - if g["code"] is None: - g["repo"] = row["repo"] - g["func_name"] = row["func_name"] - g["language"] = row["language"] - g["code"] = row["code"] - g["docstring"] = row["docstring"] - g["transformed_list"].append(row["transformed"]) - g["aug_type_list"].append(row["aug_type"]) - - # Build columnar dict, pad to max_num_augs - result: dict[str, list] = { - "repo": [], - "func_name": [], - "language": [], - "code": [], - "docstring": [], - "transformed_list": [], - "aug_type_list": [], - "num_augs": [], - "function_id": [], - } - for fid, g in groups.items(): - real_k = len(g["transformed_list"]) - if real_k == 0: - continue - truncated = g["transformed_list"][:max_num_augs] - aug_types = g["aug_type_list"][:max_num_augs] - num_augs = len(truncated) - # Pad to max_num_augs - padded_transforms = truncated + [""] * (max_num_augs - num_augs) - padded_aug_types = aug_types + [""] * (max_num_augs - num_augs) - - result["repo"].append(g["repo"]) - result["func_name"].append(g["func_name"]) - result["language"].append(g["language"]) - result["code"].append(g["code"]) - result["docstring"].append(g["docstring"]) - result["transformed_list"].append(padded_transforms) - result["aug_type_list"].append(padded_aug_types) - result["num_augs"].append(num_augs) - result["function_id"].append(fid) - - return Dataset.from_dict(result) - - -def tokenize_grouped(tokenizer, example, max_seq_length=256, max_num_augs=6): - """Tokenize a grouped example: one anchor + list of augmentations. - - Works with both single examples and batched examples (HF ``.map(batched=True)``). - """ - code_inputs = tokenizer( - example["code"], - padding="max_length", - truncation=True, - max_length=max_seq_length, - return_special_tokens_mask=True, - ) - result = { - "code_input_ids": code_inputs["input_ids"], - "code_attention_mask": code_inputs["attention_mask"], - "code_special_tokens_mask": code_inputs["special_tokens_mask"], - "function_id": example["function_id"], - "num_augs": example["num_augs"], - } - - if isinstance(example["code"], list): - # Batched mode: each element of transformed_list is a list of strings - all_aug_ids = [] - all_aug_masks = [] - all_aug_special = [] - for i, transforms in enumerate(example["transformed_list"]): - num = example["num_augs"][i] - # Only tokenize real augmentations (non-empty) - real_transforms = transforms[:num] - if real_transforms: - aug_inputs = tokenizer( - real_transforms, - padding="max_length", - truncation=True, - max_length=max_seq_length, - return_special_tokens_mask=True, - ) - all_aug_ids.append(aug_inputs["input_ids"]) - all_aug_masks.append(aug_inputs["attention_mask"]) - all_aug_special.append(aug_inputs["special_tokens_mask"]) - else: - all_aug_ids.append([]) - all_aug_masks.append([]) - all_aug_special.append([]) - result["aug_input_ids_list"] = all_aug_ids - result["aug_attention_mask_list"] = all_aug_masks - result["aug_special_tokens_mask_list"] = all_aug_special - else: - # Single mode - num = example["num_augs"] - real_transforms = example["transformed_list"][:num] - if real_transforms: - aug_inputs = tokenizer( - real_transforms, - padding="max_length", - truncation=True, - max_length=max_seq_length, - return_special_tokens_mask=True, - ) - result["aug_input_ids_list"] = aug_inputs["input_ids"] - result["aug_attention_mask_list"] = aug_inputs["attention_mask"] - result["aug_special_tokens_mask_list"] = aug_inputs["special_tokens_mask"] - else: - result["aug_input_ids_list"] = [] - result["aug_attention_mask_list"] = [] - result["aug_special_tokens_mask_list"] = [] - - return result - - def main( dataset_path: str, model_name: str, @@ -221,11 +94,11 @@ def main( sample_rate: float, checkpoint: str | None = None, tokenizer_name: str | None = None, - contra_mode: ContraMode = "info_nce", - max_num_augs: int = 6, self_contrast: bool = True, model_type: str = "roberta", pooling: str = "cls", + mlm_weight: float = 1.0, + include_nl: bool = False, ): set_seed(seed) @@ -293,31 +166,19 @@ def main( tokenizer=tokenizer, mlm=True, mlm_probability=0.15 ) - if contra_mode == "grouped": - # Regroup flat rows by function_id into {code, [aug_1, ..., aug_K]} - grouped_dataset = regroup_dataset(dataset, max_num_augs=max_num_augs) - tokenized_datasets = grouped_dataset.map( - partial( - tokenize_grouped, - tokenizer, - max_seq_length=max_seq_length, - max_num_augs=max_num_augs, - ), - batched=True, - num_proc=num_proc, - remove_columns=grouped_dataset.column_names, - ).shuffle(seed=seed) - - collator_fn = partial(grouped_contra_data_collator, mlm_collator, max_num_augs) - else: - tokenized_datasets = dataset.map( - partial(tokenize, tokenizer, max_seq_length=max_seq_length), - batched=True, - num_proc=num_proc, - remove_columns=dataset.column_names, - ).shuffle(seed=seed) + tokenized_datasets = dataset.map( + partial( + tokenize, + tokenizer, + max_seq_length=max_seq_length, + include_nl=include_nl, + ), + batched=True, + num_proc=num_proc, + remove_columns=dataset.column_names, + ).shuffle(seed=seed) - collator_fn = partial(contra_data_collator, mlm_collator) + collator_fn = partial(contra_data_collator, mlm_collator) split_dataset = tokenized_datasets.train_test_split(test_size=0.1) train_dataset = split_dataset["train"] @@ -352,8 +213,8 @@ def main( data_collator=collator_fn, processing_class=tokenizer, alpha=alpha, + mlm_weight=mlm_weight, temperature=temperature, - contra_mode=contra_mode, pooling=pooling, ) diff --git a/run_ablations.sh b/run_ablations.sh new file mode 100755 index 0000000..6c08b49 --- /dev/null +++ b/run_ablations.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +# Run all Tier 1 ablation experiments sequentially on GPUs 4,5,6,7. +set -euo pipefail + +export CUDA_VISIBLE_DEVICES=4,5,6,7 + +for cfg in experiments/ablation/*.yaml; do + echo "==========================================" + echo "Running: $cfg" + echo "Started: $(date)" + echo "==========================================" + accelerate launch --multi_gpu modeling/cli.py run "$cfg" + echo "Finished: $(date)" + echo "" +done + +echo "All ablation experiments complete." diff --git a/tests/test_config.py b/tests/test_config.py index 54eeac8..771cd20 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -4,7 +4,6 @@ import yaml from typer.testing import CliRunner -from modeling._types import ContraMode from modeling.cli import app from modeling.config import PretrainConfig, load_config @@ -36,12 +35,10 @@ def test_full_config(self) -> None: "temperature": 0.1, "max_seq_length": 128, "sample_rate": 0.5, - "contra_mode": "supcon", } ) config = load_config(path) assert config.batch_size == 32 - assert config.contra_mode == ContraMode.SUPCON assert config.resume is True assert config.seed == 42 @@ -50,7 +47,6 @@ def test_partial_config_uses_defaults(self) -> None: config = load_config(path) assert config.run_name == "minimal" assert config.batch_size == 256 # default - assert config.contra_mode == ContraMode.INFO_NCE # default def test_empty_yaml_uses_all_defaults(self) -> None: path = _write_yaml(None) @@ -61,12 +57,6 @@ def test_file_not_found(self) -> None: with pytest.raises(FileNotFoundError): load_config("/nonexistent/path.yaml") - def test_contra_mode_enum_casting(self) -> None: - for mode in ["info_nce", "supcon", "grouped"]: - path = _write_yaml({"contra_mode": mode}) - config = load_config(path) - assert isinstance(config.contra_mode, ContraMode) - class TestCli: def test_top_level_help(self) -> None: @@ -88,4 +78,3 @@ def test_pretrain_help(self) -> None: result = runner.invoke(app, ["pretrain", "--help"]) assert result.exit_code == 0 assert "--batch-size" in result.output - assert "--contra-mode" in result.output diff --git a/tests/test_grouped.py b/tests/test_grouped.py deleted file mode 100644 index 37c1fb1..0000000 --- a/tests/test_grouped.py +++ /dev/null @@ -1,293 +0,0 @@ -import torch -import torch.nn.functional as F -from transformers import DataCollatorForLanguageModeling, RobertaTokenizerFast - -from modeling.dataloader import grouped_contra_data_collator -from modeling.model import grouped_contrastive_loss, info_nce_loss -from modeling.pretrain import compute_function_id, regroup_dataset - -# --------------------------------------------------------------------------- -# grouped_contrastive_loss tests -# --------------------------------------------------------------------------- - - -class TestGroupedContrastiveLoss: - def test_identical_positives_low_loss(self): - """When all augs are copies of their anchor, loss should be low.""" - B, D, max_K = 4, 128, 2 - anchors = F.normalize(torch.randn(B, D), dim=1) - # Each anchor's augs are copies of itself - aug_list = [] - for i in range(B): - for _ in range(max_K): - aug_list.append(anchors[i]) - aug_embeddings = torch.stack(aug_list) # [B*max_K, D] - group_sizes = torch.tensor([max_K] * B) - - loss = grouped_contrastive_loss(anchors, aug_embeddings, group_sizes) - assert loss.item() >= 0.0 - assert loss.item() < 1.0 - - def test_no_augs_returns_zero(self): - """All group_sizes=0 yields loss=0.""" - B, D, max_K = 3, 64, 2 - anchors = F.normalize(torch.randn(B, D), dim=1) - aug_embeddings = torch.zeros(B * max_K, D) - group_sizes = torch.tensor([0, 0, 0]) - - loss = grouped_contrastive_loss(anchors, aug_embeddings, group_sizes) - assert loss.item() == 0.0 - - def test_gradient_flows(self): - """Loss should produce finite gradients.""" - B, D, max_K = 4, 64, 3 - anchors = torch.randn(B, D, requires_grad=True) - augs = torch.randn(B * max_K, D, requires_grad=True) - group_sizes = torch.tensor([3, 2, 1, 3]) - - loss = grouped_contrastive_loss(anchors, augs, group_sizes, temperature=0.1) - loss.backward() - assert anchors.grad is not None - assert torch.isfinite(anchors.grad).all() - assert augs.grad is not None - assert torch.isfinite(augs.grad).all() - - def test_numerical_stability_large_logits(self): - """With very large embedding values, loss should still be finite.""" - B, D, max_K = 4, 64, 2 - anchors = torch.randn(B, D) * 100 - augs = torch.randn(B * max_K, D) * 100 - group_sizes = torch.tensor([2, 2, 2, 2]) - - loss = grouped_contrastive_loss(anchors, augs, group_sizes, temperature=0.01) - assert torch.isfinite(loss) - - def test_variable_group_sizes(self): - """Mix of different aug counts per anchor.""" - B, D, max_K = 4, 64, 4 - anchors = F.normalize(torch.randn(B, D), dim=1) - augs = F.normalize(torch.randn(B * max_K, D), dim=1) - group_sizes = torch.tensor([1, 4, 2, 3]) - - loss = grouped_contrastive_loss(anchors, augs, group_sizes, temperature=0.1) - assert torch.isfinite(loss) - assert loss.item() > 0.0 - - def test_equivalent_to_infonce_when_k1(self): - """With exactly 1 aug per anchor, grouped loss ~ InfoNCE.""" - torch.manual_seed(42) - B, D = 8, 64 - query = F.normalize(torch.randn(B, D), dim=1) - key = F.normalize(torch.randn(B, D), dim=1) - - infonce = info_nce_loss(query, key, temperature=0.1) - - # Grouped: max_K=1, group_sizes all 1 - group_sizes = torch.ones(B, dtype=torch.long) - grouped = grouped_contrastive_loss(query, key, group_sizes, temperature=0.1) - - # Not exactly equal (grouped includes anchor-to-anchor terms in denom), - # but should be in the same ballpark - assert abs(infonce.item() - grouped.item()) / max(infonce.item(), 1e-6) < 1.0 - - def test_single_anchor_still_works(self): - """Edge case: B=1 should not crash (no negatives from other groups).""" - B, D, max_K = 1, 64, 3 - anchors = F.normalize(torch.randn(B, D), dim=1) - augs = F.normalize(torch.randn(B * max_K, D), dim=1) - group_sizes = torch.tensor([3]) - - loss = grouped_contrastive_loss(anchors, augs, group_sizes, temperature=0.1) - assert torch.isfinite(loss) - - def test_partial_group_padding_ignored(self): - """Padding positions (beyond group_size) should not affect loss.""" - torch.manual_seed(123) - B, D = 3, 64 - anchors = F.normalize(torch.randn(B, D), dim=1) - - # max_K=3: group 0 has 2 real augs, group 1 has 1, group 2 has 3 - max_K = 3 - augs = F.normalize(torch.randn(B * max_K, D), dim=1) - group_sizes = torch.tensor([2, 1, 3]) - loss1 = grouped_contrastive_loss(anchors, augs, group_sizes, temperature=0.1) - - # Change padding positions (index 2 for group 0, indices 4-5 for group 1) - augs_modified = augs.clone() - augs_modified[2] = torch.randn(D) # group 0, slot 2 (padding) - augs_modified[4] = torch.randn(D) # group 1, slot 1 (padding) - augs_modified[5] = torch.randn(D) # group 1, slot 2 (padding) - loss2 = grouped_contrastive_loss( - anchors, augs_modified, group_sizes, temperature=0.1 - ) - - assert torch.isclose(loss1, loss2, atol=1e-6) - - -# --------------------------------------------------------------------------- -# grouped_contra_data_collator tests -# --------------------------------------------------------------------------- - - -def _make_grouped_feature(seq_len, aug_counts): - """Create a mock grouped feature dict for testing the collator.""" - features = [] - for n_augs in aug_counts: - f = { - "code_input_ids": [1] + [100] * (seq_len - 2) + [2], - "code_attention_mask": [1] * seq_len, - "code_special_tokens_mask": [1] + [0] * (seq_len - 2) + [1], - "aug_input_ids_list": [], - "aug_attention_mask_list": [], - "aug_special_tokens_mask_list": [], - "function_id": hash(str(len(features))) & 0x7FFFFFFFFFFFFFFF, - } - for k in range(n_augs): - f["aug_input_ids_list"].append([1] + [200 + k] * (seq_len - 2) + [2]) - f["aug_attention_mask_list"].append([1] * seq_len) - f["aug_special_tokens_mask_list"].append([1] + [0] * (seq_len - 2) + [1]) - features.append(f) - return features - - -class TestGroupedContraDataCollator: - def _get_mlm_collator(self): - tokenizer = RobertaTokenizerFast.from_pretrained("microsoft/codebert-base") - return DataCollatorForLanguageModeling( - tokenizer=tokenizer, mlm=True, mlm_probability=0.15 - ) - - def test_output_shapes(self): - """Verify correct tensor shapes for code and aug batches.""" - mlm_collator = self._get_mlm_collator() - seq_len = 16 - features = _make_grouped_feature(seq_len, aug_counts=[2, 3]) - batch = grouped_contra_data_collator(mlm_collator, 6, features) - - B = 2 - max_K = 3 # max(2, 3) - assert batch["code_input_ids"].shape == (B, seq_len) - assert batch["aug_input_ids"].shape == (B * max_K, seq_len) - assert batch["group_sizes"].shape == (B,) - assert batch["group_sizes"].tolist() == [2, 3] - - def test_padding_has_no_mlm_labels(self): - """Padding aug slots should have all labels=-100 (no MLM masking).""" - mlm_collator = self._get_mlm_collator() - seq_len = 16 - # Group 0: 1 aug, Group 1: 3 augs → max_K=3, group 0 has 2 padding slots - features = _make_grouped_feature(seq_len, aug_counts=[1, 3]) - batch = grouped_contra_data_collator(mlm_collator, 6, features) - - # Group 0's padding slots are indices 1 and 2 in the flattened aug batch - # (group 0 occupies slots 0..2, real=1, padding=slots 1,2) - padding_labels_1 = batch["aug_labels"][1] - padding_labels_2 = batch["aug_labels"][2] - assert (padding_labels_1 == -100).all() - assert (padding_labels_2 == -100).all() - - def test_group_sizes_correct(self): - """group_sizes should reflect actual aug counts.""" - mlm_collator = self._get_mlm_collator() - features = _make_grouped_feature(16, aug_counts=[1, 2, 4]) - batch = grouped_contra_data_collator(mlm_collator, 6, features) - assert batch["group_sizes"].tolist() == [1, 2, 4] - - def test_max_num_augs_truncation(self): - """Features with more augs than max_num_augs get truncated.""" - mlm_collator = self._get_mlm_collator() - features = _make_grouped_feature(16, aug_counts=[5, 3]) - batch = grouped_contra_data_collator(mlm_collator, 2, features) - - B = 2 - max_K = 2 - assert batch["aug_input_ids"].shape == (B * max_K, 16) - assert batch["group_sizes"].tolist() == [2, 2] - - def test_function_id_passed_through(self): - """function_id should be present in the batch.""" - mlm_collator = self._get_mlm_collator() - features = _make_grouped_feature(16, aug_counts=[2, 1]) - batch = grouped_contra_data_collator(mlm_collator, 6, features) - assert "function_id" in batch - assert batch["function_id"].shape == (2,) - - -# --------------------------------------------------------------------------- -# regroup_dataset tests -# --------------------------------------------------------------------------- - - -class TestRegroupDataset: - def _make_flat_dataset(self): - """Create a mock flat HF dataset with multiple rows per function.""" - from datasets import Dataset - - data = { - "repo": ["r1", "r1", "r1", "r2", "r2"], - "func_name": ["f1", "f1", "f1", "f2", "f2"], - "language": ["python"] * 5, - "code": [ - "def foo(): pass", - "def foo(): pass", - "def foo(): pass", - "def bar(): pass", - "def bar(): pass", - ], - "docstring": [""] * 5, - "transformed": [ - "def foo_v1(): pass", - "def foo_v2(): pass", - "def foo_v3(): pass", - "def bar_v1(): pass", - "def bar_v2(): pass", - ], - "aug_type": [ - "LocalVarRenaming", - "ReverseIfElse", - "AddAssignment2EqualAssignment", - "LocalVarRenaming", - "ReverseIfElse", - ], - } - return Dataset.from_dict(data) - - def test_groups_by_code(self): - """Rows with the same code should be grouped together.""" - dataset = self._make_flat_dataset() - grouped = regroup_dataset(dataset, max_num_augs=6) - assert len(grouped) == 2 # two distinct functions - - def test_preserves_all_augmentations(self): - """All augmentations should be preserved in the grouped output.""" - dataset = self._make_flat_dataset() - grouped = regroup_dataset(dataset, max_num_augs=6) - - # Find the group with 3 augs (foo) and 2 augs (bar) - total_real_augs = sum(grouped["num_augs"]) - assert total_real_augs == 5 # 3 + 2 - - def test_padding_to_max_num_augs(self): - """transformed_list should be padded to max_num_augs with empty strings.""" - dataset = self._make_flat_dataset() - grouped = regroup_dataset(dataset, max_num_augs=4) - - for transforms in grouped["transformed_list"]: - assert len(transforms) == 4 - - def test_truncation(self): - """Groups with more augs than max_num_augs get truncated.""" - dataset = self._make_flat_dataset() - grouped = regroup_dataset(dataset, max_num_augs=2) - - for num in grouped["num_augs"]: - assert num <= 2 - - def test_function_id_deterministic(self): - """function_id should match compute_function_id on the code string.""" - dataset = self._make_flat_dataset() - grouped = regroup_dataset(dataset, max_num_augs=6) - - for i in range(len(grouped)): - expected_fid = compute_function_id(grouped["code"][i]) - assert grouped["function_id"][i] == expected_fid diff --git a/tests/test_supcon.py b/tests/test_supcon.py index f2e7556..beeef9c 100644 --- a/tests/test_supcon.py +++ b/tests/test_supcon.py @@ -1,7 +1,7 @@ import torch import torch.nn.functional as F -from modeling.model import build_positive_mask, info_nce_loss, supcon_loss +from modeling.model import build_positive_mask, supcon_loss class TestBuildPositiveMask: @@ -80,20 +80,3 @@ def test_batch_with_mixed_positives(self): loss = supcon_loss(embeddings, ids, temperature=0.1) assert torch.isfinite(loss) assert loss.item() > 0.0 - - def test_supcon_similar_to_infonce_single_positive(self): - """With exactly 1 positive per anchor (code/aug pairs), SupCon and - InfoNCE should produce comparable losses.""" - torch.manual_seed(42) - query = F.normalize(torch.randn(4, 64), dim=1) - key = F.normalize(torch.randn(4, 64), dim=1) - - infonce = info_nce_loss(query, key, temperature=0.1) - - # SupCon: concatenate, ids are [0,1,2,3,0,1,2,3] - all_emb = torch.cat([query, key], dim=0) - ids = torch.cat([torch.arange(4), torch.arange(4)]) - sc = supcon_loss(all_emb, ids, temperature=0.1) - - # Not exactly equal (different denominator sizes), but same ballpark - assert abs(infonce.item() - sc.item()) / max(infonce.item(), 1e-6) < 1.0