From 7a535703552fcdae91cf253d32042e767dfd0c0f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 12 Mar 2026 02:34:28 +0000 Subject: [PATCH 1/4] Initial plan From 13f329b9285d0aedbb20c492f289c2a57d406ee5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 12 Mar 2026 02:54:41 +0000 Subject: [PATCH 2/4] Implement dynamic approaches for code simplification (dataset generator, models, training, integration, evaluation)" Co-authored-by: hanyiliu <35943468+hanyiliu@users.noreply.github.com> --- src/data/dynamic/__init__.py | 0 src/data/dynamic/dynamic_dataset_generator.py | 540 ++++++++++++++++++ src/evaluation/dynamic/__init__.py | 0 src/evaluation/dynamic/dynamic_evaluator.py | 514 +++++++++++++++++ src/models/dynamic/__init__.py | 0 src/models/dynamic/cnn_model.py | 104 ++++ src/models/dynamic/linear_model.py | 121 ++++ src/models/dynamic/rnn_model.py | 121 ++++ src/models/integrated/__init__.py | 0 src/models/integrated/integrated_codebert.py | 274 +++++++++ src/models/integrated/integrated_codet5.py | 249 ++++++++ src/training/dynamic/__init__.py | 0 src/training/dynamic/train_dynamic.py | 452 +++++++++++++++ 13 files changed, 2375 insertions(+) create mode 100644 src/data/dynamic/__init__.py create mode 100644 src/data/dynamic/dynamic_dataset_generator.py create mode 100644 src/evaluation/dynamic/__init__.py create mode 100644 src/evaluation/dynamic/dynamic_evaluator.py create mode 100644 src/models/dynamic/__init__.py create mode 100644 src/models/dynamic/cnn_model.py create mode 100644 src/models/dynamic/linear_model.py create mode 100644 src/models/dynamic/rnn_model.py create mode 100644 src/models/integrated/__init__.py create mode 100644 src/models/integrated/integrated_codebert.py create mode 100644 src/models/integrated/integrated_codet5.py create mode 100644 src/training/dynamic/__init__.py create mode 100644 src/training/dynamic/train_dynamic.py diff --git a/src/data/dynamic/__init__.py b/src/data/dynamic/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/data/dynamic/dynamic_dataset_generator.py b/src/data/dynamic/dynamic_dataset_generator.py new file mode 100644 index 0000000..1837322 --- /dev/null +++ b/src/data/dynamic/dynamic_dataset_generator.py @@ -0,0 +1,540 @@ +""" +Dynamic Dataset Generator +========================== +Generates (code snippet, per-token attention score) training pairs from +fine-tuned CodeBERT and CodeT5 checkpoints in ``checkpoints/``. + +For each code snippet in the CodeSearchNet dataset the generator: + +1. Tokenises the snippet with the model's tokenizer. +2. Passes it through a fine-tuned model with attention outputs enabled. +3. Extracts per-token importance scores: + - CodeBERT (code search): CLS attention from the last encoder layer, + averaged across all heads. + - CodeT5 (summarization): encoder-decoder cross-attention from the last + decoder layer, max-pooled over decoder time-steps. +4. Aligns each subword token to its AST statement category. +5. Stores a serialisable record: + {idx, code, token_ids, positions, category_ids, attention_scores} + +The resulting dataset is consumed by the dynamic model trainers in +``src/training/dynamic/``. + +Output +------ +``data/dynamic/dynamic_{model}_{task}.json`` + +Each record has the form:: + + { + "idx": int, + "code": str, + "token_ids": list[int], + "positions": list[int], + "category_ids": list[int], + "attention_scores": list[float] + } + +Usage +----- + python -m src.data.dynamic.dynamic_dataset_generator \\ + --model codebert --task search \\ + --checkpoint checkpoints/codebert_encoder_only/epoch_1 \\ + --output-dir data/dynamic --demo 500 + + python -m src.data.dynamic.dynamic_dataset_generator \\ + --model codet5 --task summarization \\ + --checkpoint checkpoints/codet5_encoder_decoder/epoch_1 \\ + --output-dir data/dynamic +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +from typing import cast + +import torch +from torch.utils.data import DataLoader +from tqdm import tqdm +from transformers.models.roberta import RobertaTokenizer +from transformers.models.roberta.modeling_roberta import RobertaModel +from transformers.models.t5 import T5ForConditionalGeneration + +from src.data.data_loader import load_codesearchnet +from src.attention.attention_extraction import ( + map_subword_tokens_to_categories, + _collate, + _cls_scores_from_outputs, + SPECIAL_TOKEN_STRINGS, +) + +# --------------------------------------------------------------------------- +# Category vocabulary (statement-level, matching attention_extraction.py) +# --------------------------------------------------------------------------- + +ALL_CATEGORIES: list[str] = [ + "If Condition", + "For Loop", + "While Loop", + "Switch", + "Case", + "Break", + "Continue", + "Try", + "Catch", + "Finally", + "Throw", + "Variable Declaration", + "Field Declaration", + "Assignment", + "Expression", + "Update Expression", + "Method Signature", + "Constructor", + "Method Call", + "Object Creation", + "Return", + "Synchronized", + "Annotation", + "Getter Method", + "Setter Method", + "Method Declaration", + "Logging Call", + "Other", +] + +CATEGORY_TO_ID: dict[str, int] = {cat: i for i, cat in enumerate(ALL_CATEGORIES)} +NUM_CATEGORIES: int = len(ALL_CATEGORIES) + +MAX_LENGTH = 512 +BATCH_SIZE = 32 + + +# --------------------------------------------------------------------------- +# Model loading (from fine-tuned checkpoint) +# --------------------------------------------------------------------------- + +def load_finetuned_codebert( + checkpoint_dir: str, + device: torch.device, +) -> tuple[RobertaTokenizer, RobertaModel]: + """Load a fine-tuned CodeBERT checkpoint with attention outputs enabled. + + Args: + checkpoint_dir: Path produced by ``FinetunedCodeBERT.save_model()``. + device: Target device. + + Returns: + (tokenizer, model) — model is in eval mode with attention outputs. + """ + print(f"Loading fine-tuned CodeBERT from '{checkpoint_dir}'...") + tokenizer = cast( + RobertaTokenizer, + RobertaTokenizer.from_pretrained(checkpoint_dir), + ) + model = cast( + RobertaModel, + RobertaModel.from_pretrained(checkpoint_dir, output_attentions=True), + ) + cast(torch.nn.Module, model).to(device).eval() + return tokenizer, model + + +def load_finetuned_codet5( + checkpoint_dir: str, + device: torch.device, +) -> tuple[RobertaTokenizer, T5ForConditionalGeneration]: + """Load a fine-tuned CodeT5 checkpoint with attention outputs enabled. + + Falls back gracefully if the tokenizer metadata uses an unexpected + ``extra_special_tokens`` field (older HF versions). + + Args: + checkpoint_dir: Path produced by ``FineTunedCodeT5.save_model()``. + device: Target device. + + Returns: + (tokenizer, model) — model is in eval mode with attention outputs. + """ + print(f"Loading fine-tuned CodeT5 from '{checkpoint_dir}'...") + try: + tokenizer = cast( + RobertaTokenizer, + RobertaTokenizer.from_pretrained(checkpoint_dir), + ) + except TypeError as exc: + if "extra_special_tokens" not in str(exc): + raise + tokenizer = cast( + RobertaTokenizer, + RobertaTokenizer.from_pretrained(checkpoint_dir, extra_special_tokens=[]), + ) + model = cast( + T5ForConditionalGeneration, + T5ForConditionalGeneration.from_pretrained(checkpoint_dir, output_attentions=True), + ) + cast(torch.nn.Module, model).to(device).eval() + return tokenizer, model + + +# --------------------------------------------------------------------------- +# Batched score extraction +# --------------------------------------------------------------------------- + +@torch.no_grad() +def _extract_cls_scores( + batch: list[dict], + tokenizer: RobertaTokenizer, + model: RobertaModel, + device: torch.device, +) -> list[tuple[list[int], list[int], list[int], list[float]]]: + """Extract CLS attention scores for a batch (CodeBERT code-search). + + Returns a list of (token_ids, positions, category_ids, scores) per sample. + Padding tokens are stripped so all lists reflect only the real sequence. + """ + codes = [s.get("func_code_string", "") for s in batch] + enc = tokenizer( + codes, + return_tensors="pt", + truncation=True, + max_length=MAX_LENGTH, + padding=True, + ).to(device) + + outputs = model(**enc, output_attentions=True) + raw_scores = _cls_scores_from_outputs(outputs, len(batch)) + + results = [] + for i, code in enumerate(codes): + real_len = int(enc["attention_mask"][i].sum().item()) + ids_tensor = enc["input_ids"][i, :real_len] + token_ids = ids_tensor.cpu().tolist() + positions = list(range(real_len)) + scores = raw_scores[i][:real_len] + + tokens = tokenizer.convert_ids_to_tokens(ids_tensor) + categories = map_subword_tokens_to_categories(code, tokens) + category_ids = [ + CATEGORY_TO_ID.get(cat, CATEGORY_TO_ID["Other"]) + for cat in categories + ] + results.append((token_ids, positions, category_ids, scores)) + return results + + +@torch.no_grad() +def _extract_cross_attn_scores( + batch: list[dict], + tokenizer: RobertaTokenizer, + model: T5ForConditionalGeneration, + device: torch.device, +) -> list[tuple[list[int], list[int], list[int], list[float]]]: + """Extract encoder-decoder cross-attention scores (CodeT5 summarization). + + Uses teacher-forcing with the reference description as the decoder input, + matching the approach in ``attention_extraction.extract_batch_enc_dec``. + Returns a list of (token_ids, positions, category_ids, scores) per sample. + """ + codes = [s.get("func_code_string", "") for s in batch] + descriptions = [s.get("func_documentation_string", "") for s in batch] + + enc = tokenizer( + codes, + return_tensors="pt", + truncation=True, + max_length=MAX_LENGTH, + padding=True, + ).to(device) + + has_desc = any(descriptions) + if has_desc: + dec_enc = tokenizer( + descriptions, + return_tensors="pt", + truncation=True, + max_length=128, + padding=True, + ).to(device) + decoder_input_ids = dec_enc["input_ids"] + else: + bos_id = tokenizer.pad_token_id or tokenizer.eos_token_id or 0 + decoder_input_ids = torch.full((len(batch), 1), bos_id, device=device) + + outputs = model( + input_ids=enc["input_ids"], + attention_mask=enc["attention_mask"], + decoder_input_ids=decoder_input_ids, + output_attentions=True, + ) + + # cross_attentions[-1]: [B, heads, tgt, src] → avg heads → max over tgt + cross_avg = outputs.cross_attentions[-1].mean(dim=1) # [B, tgt, src] + max_scores = cross_avg.max(dim=1).values # [B, src] + + results = [] + for i, code in enumerate(codes): + real_len = int(enc["attention_mask"][i].sum().item()) + ids_tensor = enc["input_ids"][i, :real_len] + token_ids = ids_tensor.cpu().tolist() + positions = list(range(real_len)) + scores = max_scores[i, :real_len].cpu().tolist() + + tokens = tokenizer.convert_ids_to_tokens(ids_tensor) + categories = map_subword_tokens_to_categories(code, tokens) + category_ids = [ + CATEGORY_TO_ID.get(cat, CATEGORY_TO_ID["Other"]) + for cat in categories + ] + results.append((token_ids, positions, category_ids, scores)) + return results + + +# --------------------------------------------------------------------------- +# Dataset generation entry point +# --------------------------------------------------------------------------- + +def generate_dataset( + model_type: str, + task: str, + checkpoint_dir: str, + output_dir: str = "data/dynamic", + language: str = "java", + split: str = "train", + batch_size: int = BATCH_SIZE, + demo: int | None = None, + device_str: str | None = None, +) -> str: + """Generate a dynamic training dataset and save it to disk. + + Args: + model_type: ``'codebert'`` or ``'codet5'``. + task: ``'search'`` or ``'summarization'``. + checkpoint_dir: Path to a fine-tuned model checkpoint directory. + output_dir: Directory to write the output JSON file. + language: CodeSearchNet language (default ``'java'``). + split: Dataset split to process (default ``'train'``). + batch_size: Samples per batch during extraction. + demo: If set, process only the first *demo* samples. + device_str: Force a specific device (e.g. ``'cpu'``). + + Returns: + Path to the saved JSON file. + + Raises: + ValueError: For invalid ``model_type`` / ``task`` combinations. + """ + if model_type == "codebert" and task == "summarization": + raise ValueError( + "CodeBERT has no decoder. Use --model codet5 --task summarization." + ) + + if device_str: + device = torch.device(device_str) + elif torch.cuda.is_available(): + device = torch.device("cuda") + elif torch.backends.mps.is_available(): + device = torch.device("mps") + else: + device = torch.device("cpu") + + print( + f"Dynamic dataset generator | model={model_type} | task={task} | " + f"checkpoint={checkpoint_dir} | device={device}" + ) + + # Load dataset + dataset = load_codesearchnet(language=language)[split] + if demo is not None: + dataset = dataset.select(range(demo)) + print(f"Samples: {len(dataset):,}") + + # Load fine-tuned model + if model_type == "codebert": + tokenizer, model = load_finetuned_codebert(checkpoint_dir, device) + extractor = "cls" + else: + tokenizer, model = load_finetuned_codet5(checkpoint_dir, device) + extractor = "cross_attn" + + loader = DataLoader( + dataset, + batch_size=batch_size, + collate_fn=_collate, + num_workers=0, + ) + + records: list[dict] = [] + errors = 0 + sample_idx = 0 + + for batch in tqdm(loader, desc="Generating"): + try: + if extractor == "cls": + batch_out = _extract_cls_scores( + batch, + tokenizer, + cast(RobertaModel, model), + device, + ) + else: + batch_out = _extract_cross_attn_scores( + batch, + tokenizer, + cast(T5ForConditionalGeneration, model), + device, + ) + except Exception as exc: + errors += len(batch) + if errors <= 10: + print(f"\nBatch error at sample ~{sample_idx}: {exc}") + sample_idx += len(batch) + continue + + for sample, (token_ids, positions, category_ids, scores) in zip(batch, batch_out): + records.append( + { + "idx": sample_idx, + "code": sample.get("func_code_string", ""), + "token_ids": token_ids, + "positions": positions, + "category_ids": category_ids, + "attention_scores": scores, + } + ) + sample_idx += 1 + + print(f"Done — {len(records):,} records, {errors} errors.") + + Path(output_dir).mkdir(parents=True, exist_ok=True) + out_path = os.path.join(output_dir, f"dynamic_{model_type}_{task}.json") + with open(out_path, "w", encoding="utf-8") as f: + json.dump(records, f) + print(f"Saved → {out_path}") + return out_path + + +# --------------------------------------------------------------------------- +# PyTorch Dataset wrapper for the generated data +# --------------------------------------------------------------------------- + +class DynamicScoringDataset(torch.utils.data.Dataset): + """PyTorch Dataset for training dynamic scoring models. + + Loads the JSON file produced by :func:`generate_dataset` and serves + one record at a time as a dict with the following keys: + + - ``token_ids`` : ``LongTensor`` [seq_len] + - ``positions`` : ``LongTensor`` [seq_len] + - ``category_ids`` : ``LongTensor`` [seq_len] + - ``attention_scores`` : ``FloatTensor`` [seq_len] + + Usage:: + + ds = DynamicScoringDataset("data/dynamic/dynamic_codebert_search.json") + loader = DataLoader(ds, batch_size=32, collate_fn=dynamic_collate_fn) + """ + + def __init__(self, json_path: str) -> None: + with open(json_path, encoding="utf-8") as f: + self._records: list[dict] = json.load(f) + + def __len__(self) -> int: + return len(self._records) + + def __getitem__(self, idx: int) -> dict: + rec = self._records[idx] + return { + "token_ids": torch.tensor(rec["token_ids"], dtype=torch.long), + "positions": torch.tensor(rec["positions"], dtype=torch.long), + "category_ids": torch.tensor(rec["category_ids"], dtype=torch.long), + "attention_scores": torch.tensor(rec["attention_scores"], dtype=torch.float), + } + + +def dynamic_collate_fn(batch: list[dict]) -> dict: + """Pad a batch of variable-length sequences for DataLoader. + + Pads ``token_ids``, ``positions``, ``category_ids``, and + ``attention_scores`` to the maximum sequence length in the batch. + A boolean ``padding_mask`` tensor is included so models can ignore + padding positions when computing loss. + + Returns: + Dict with keys: + - ``token_ids`` : ``LongTensor`` [B, max_len] + - ``positions`` : ``LongTensor`` [B, max_len] + - ``category_ids`` : ``LongTensor`` [B, max_len] + - ``attention_scores`` : ``FloatTensor`` [B, max_len] + - ``lengths`` : ``LongTensor`` [B] + - ``padding_mask`` : ``BoolTensor`` [B, max_len] True = valid + """ + lengths = torch.tensor([item["token_ids"].size(0) for item in batch], dtype=torch.long) + max_len = int(lengths.max().item()) + B = len(batch) + + token_ids = torch.zeros(B, max_len, dtype=torch.long) + positions = torch.zeros(B, max_len, dtype=torch.long) + category_ids = torch.zeros(B, max_len, dtype=torch.long) + attention_scores = torch.zeros(B, max_len, dtype=torch.float) + padding_mask = torch.zeros(B, max_len, dtype=torch.bool) + + for i, (item, L) in enumerate(zip(batch, lengths.tolist())): + token_ids[i, :L] = item["token_ids"] + positions[i, :L] = item["positions"] + category_ids[i, :L] = item["category_ids"] + attention_scores[i, :L] = item["attention_scores"] + padding_mask[i, :L] = True + + return { + "token_ids": token_ids, + "positions": positions, + "category_ids": category_ids, + "attention_scores": attention_scores, + "lengths": lengths, + "padding_mask": padding_mask, + } + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description="Dynamic Dataset Generator") + p.add_argument("--model", choices=["codebert", "codet5"], default="codebert") + p.add_argument("--task", choices=["search", "summarization"], default="search") + p.add_argument("--checkpoint", required=True, + help="Path to fine-tuned model checkpoint directory.") + p.add_argument("--output-dir", default="data/dynamic") + p.add_argument("--language", default="java") + p.add_argument("--split", default="train", + choices=["train", "validation", "test"]) + p.add_argument("--batch-size", type=int, default=BATCH_SIZE) + p.add_argument("--demo", type=int, default=None, metavar="N", + help="Process only the first N samples.") + p.add_argument("--device", default=None) + return p.parse_args() + + +def main() -> None: + """CLI entry point.""" + args = _parse_args() + generate_dataset( + model_type = args.model, + task = args.task, + checkpoint_dir = args.checkpoint, + output_dir = args.output_dir, + language = args.language, + split = args.split, + batch_size = args.batch_size, + demo = args.demo, + device_str = args.device, + ) + + +if __name__ == "__main__": + main() diff --git a/src/evaluation/dynamic/__init__.py b/src/evaluation/dynamic/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/evaluation/dynamic/dynamic_evaluator.py b/src/evaluation/dynamic/dynamic_evaluator.py new file mode 100644 index 0000000..1252c33 --- /dev/null +++ b/src/evaluation/dynamic/dynamic_evaluator.py @@ -0,0 +1,514 @@ +""" +Dynamic Model Evaluator +======================== +Computes MRR (code search) and BLEU-4 (code summarisation) for the integrated +dynamic models: ``IntegratedCodeBERT`` and ``IntegratedCodeT5``. + +The evaluator mirrors the API of ``CodeSearchEvaluator`` and +``CodeSummarizationEvaluator`` from ``src/evaluation/``, but wraps the +integrated models so that each forward call first simplifies the code with the +lightweight dynamic scoring model before running the downstream transformer. + +Usage +----- +CLI:: + + # MRR with CodeBERT + dynamic RNN scorer (30% removal) + python -m src.evaluation.dynamic.dynamic_evaluator \\ + --task search \\ + --model codebert \\ + --dynamic-checkpoint checkpoints/dynamic/rnn_codebert_search/best.pt \\ + --downstream-checkpoint checkpoints/codebert_encoder_only/epoch_1 \\ + --simplification-ratio 0.30 + + # BLEU-4 with CodeT5 + dynamic CNN scorer (20% removal) + python -m src.evaluation.dynamic.dynamic_evaluator \\ + --task summarization \\ + --model codet5 \\ + --dynamic-checkpoint checkpoints/dynamic/cnn_codet5_summarization/best.pt \\ + --downstream-checkpoint checkpoints/codet5_encoder_decoder/epoch_1 \\ + --simplification-ratio 0.20 + +Programmatic:: + + from src.evaluation.dynamic.dynamic_evaluator import DynamicEvaluator + + evaluator = DynamicEvaluator( + task = "search", + model_type = "codebert", + dynamic_checkpoint = "checkpoints/dynamic/rnn_codebert_search/best.pt", + downstream_checkpoint = "checkpoints/codebert_encoder_only/epoch_1", + simplification_ratio = 0.30, + ) + result = evaluator.evaluate(test_dataset) + print(result) # {"mrr": 0.xxxx, "elapsed": xx.x} + + evaluator_summ = DynamicEvaluator( + task = "summarization", + model_type = "codet5", + dynamic_checkpoint = "checkpoints/dynamic/cnn_codet5_summ/best.pt", + downstream_checkpoint = "checkpoints/codet5_encoder_decoder/epoch_1", + simplification_ratio = 0.20, + ) + result = evaluator_summ.evaluate(test_dataset) + print(result) # {"bleu4": xx.xx, "elapsed": xx.x} +""" + +from __future__ import annotations + +import argparse +import json +import random +import time +from pathlib import Path +from typing import Any, Optional + +import torch +from nltk.translate.bleu_score import SmoothingFunction, corpus_bleu +from tqdm import tqdm + +from src.data.data_loader import load_codesearchnet +from src.models.integrated.integrated_codebert import IntegratedCodeBERT +from src.models.integrated.integrated_codet5 import IntegratedCodeT5 + +# ───────────────────────────────────────────────────────────────────────────── +# Constants +# ───────────────────────────────────────────────────────────────────────────── + +MAX_LENGTH = 512 +POOL_SIZE = 1000 +EVAL_SEED = 42 +RATIOS = [0.10, 0.20, 0.30, 0.40, 0.50] + + +# ───────────────────────────────────────────────────────────────────────────── +# Dataset normalisation helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _iter_samples(dataset) -> list[dict]: + """Normalise a HuggingFace dataset or JSON list to flat dicts. + + Each dict has keys ``'code'`` and ``'query'`` / ``'reference'``. + """ + samples: list[dict] = [] + for item in dataset: + if isinstance(item, dict): + code = item.get("simplified_code") or item.get("func_code_string", "") + query = item.get("description") or item.get("func_documentation_string", "") + else: + code = item["func_code_string"] + query = item["func_documentation_string"] + if code.strip() and query.strip(): + samples.append({"code": code, "query": query, "reference": query}) + return samples + + +# ───────────────────────────────────────────────────────────────────────────── +# DynamicEvaluator +# ───────────────────────────────────────────────────────────────────────────── + +class DynamicEvaluator: + """Evaluates an integrated dynamic model on code search (MRR) or + code summarisation (BLEU-4). + + Args: + task: ``'search'`` or ``'summarization'``. + model_type: ``'codebert'`` or ``'codet5'``. + dynamic_checkpoint: Path to the ``.pt`` file produced by + ``train_dynamic.py``. + downstream_checkpoint: Path to a fine-tuned CodeBERT / CodeT5 + checkpoint. Required for search; optional for + summarisation (falls back to pretrained weights). + simplification_ratio: Fraction of tokens to remove dynamically. + device: Device string; auto-detected if ``None``. + batch_size: Pair-scoring batch size (search only). + + Usage:: + + evaluator = DynamicEvaluator( + task="search", model_type="codebert", + dynamic_checkpoint="...", downstream_checkpoint="...", + ) + results = evaluator.evaluate(test_dataset) + """ + + def __init__( + self, + task: str, + model_type: str = "codebert", + dynamic_checkpoint: str = "", + downstream_checkpoint: Optional[str] = None, + simplification_ratio: float = 0.30, + device: Optional[str] = None, + batch_size: int = 64, + ) -> None: + if task not in ("search", "summarization"): + raise ValueError(f"task must be 'search' or 'summarization', got '{task}'") + if model_type not in ("codebert", "codet5"): + raise ValueError(f"model_type must be 'codebert' or 'codet5', got '{model_type}'") + if task == "summarization" and model_type == "codebert": + raise ValueError("CodeBERT has no decoder. Use model_type='codet5' for summarization.") + if task == "search" and not downstream_checkpoint: + raise ValueError("Search evaluation requires downstream_checkpoint with classifier head.") + + self.task = task + self.model_type = model_type + self.simplification_ratio = simplification_ratio + self.batch_size = batch_size + + if device: + self.device = torch.device(device) + elif torch.cuda.is_available(): + self.device = torch.device("cuda") + elif torch.backends.mps.is_available(): + self.device = torch.device("mps") + else: + self.device = torch.device("cpu") + + print( + f"[DynamicEvaluator] task={task} | model={model_type} | " + f"ratio={simplification_ratio:.0%} | device={self.device}" + ) + + if model_type == "codebert": + self._model: Any = IntegratedCodeBERT( + dynamic_checkpoint = dynamic_checkpoint, + codebert_checkpoint = downstream_checkpoint, + simplification_ratio = simplification_ratio, + device = device, + ) + else: + self._model = IntegratedCodeT5( + dynamic_checkpoint = dynamic_checkpoint, + codet5_checkpoint = downstream_checkpoint, + simplification_ratio = simplification_ratio, + device = device, + ) + + _, tokenizer = self._model._load_model() + self._tokenizer: Any = tokenizer + + # ── Code-search MRR ────────────────────────────────────────────────────── + + @torch.no_grad() + def _score_pairs( + self, + query: str, + codes: list[str], + ) -> torch.Tensor: + """Score (query, code) pairs with class-1 logit after simplification.""" + all_scores: list[torch.Tensor] = [] + + # Simplify each code snippet before scoring + simplified_codes = [self._model.simplify_code(c) or c for c in codes] + + for start in range(0, len(simplified_codes), self.batch_size): + chunk = simplified_codes[start : start + self.batch_size] + query_chunk = [query] * len(chunk) + enc = self._tokenizer( + query_chunk, + chunk, + return_tensors = "pt", + truncation = True, + max_length = MAX_LENGTH, + padding = True, + ).to(self.device) + + logits = self._model.forward_classification( + input_ids = enc["input_ids"], + attention_mask = enc["attention_mask"], + ) + if isinstance(logits, tuple): + logits = logits[0] + all_scores.append(logits[:, 1].detach().cpu()) + + return torch.cat(all_scores, dim=0) + + def compute_mrr( + self, + dataset, + pool_size: int = POOL_SIZE, + max_samples: Optional[int] = None, + seed: int = EVAL_SEED, + ) -> tuple[float, float]: + """Compute MRR on dataset with dynamic simplification. + + Args: + dataset: HuggingFace Dataset or list[dict]. + pool_size: Candidate pool per query (1 correct + negatives). + max_samples: Evaluate on at most this many queries. + seed: RNG seed for reproducible negative sampling. + + Returns: + ``(mrr, elapsed_seconds)`` + """ + rng = random.Random(seed) + samples = _iter_samples(dataset) + if max_samples is not None: + samples = samples[:max_samples] + n = len(samples) + if n == 0: + return 0.0, 0.0 + + t_start = time.time() + reciprocal_ranks = [] + + for i, sample in enumerate(tqdm(samples, desc="MRR (dynamic)")): + query = sample["query"] + neg_indices = rng.sample( + [j for j in range(n) if j != i], + k=min(pool_size - 1, n - 1), + ) + pool_indices = [i] + neg_indices + pool_codes = [samples[j]["code"] for j in pool_indices] + + scores = self._score_pairs(query, pool_codes) + ranked = scores.argsort(descending=True).tolist() + rank_of_correct = ranked.index(0) + 1 + reciprocal_ranks.append(1.0 / rank_of_correct) + + elapsed = time.time() - t_start + mrr = sum(reciprocal_ranks) / len(reciprocal_ranks) + return mrr, elapsed + + # ── Code-summarisation BLEU-4 ──────────────────────────────────────────── + + def compute_bleu( + self, + dataset, + max_samples: Optional[int] = None, + log_every: int = 100, + ) -> tuple[float, float]: + """Compute corpus BLEU-4 on dataset with dynamic simplification. + + Args: + dataset: HuggingFace Dataset or list[dict]. + max_samples: Evaluate on at most this many examples. + log_every: Print progress every N examples. + + Returns: + ``(bleu4, elapsed_seconds)`` + """ + samples = _iter_samples(dataset) + if max_samples is not None: + samples = samples[:max_samples] + n = len(samples) + if n == 0: + print("[DynamicEvaluator] Empty dataset — returning 0.") + return 0.0, 0.0 + + print(f"\n[DynamicEvaluator] BLEU-4 evaluation on {n:,} samples...") + + hypotheses: list[list[str]] = [] + references: list[list[list[str]]] = [] + t_start = time.time() + + for i, sample in enumerate(samples): + try: + pred = self._model.forward_generation(sample["code"]) + except Exception as exc: + print(f" [WARNING] sample {i} failed: {exc}") + pred = "" + + hypotheses.append(pred.split()) + references.append([sample["reference"].split()]) + + if (i + 1) % log_every == 0 or (i + 1) == n: + print(f" Progress: {i + 1:,}/{n:,}") + + elapsed = time.time() - t_start + sf = SmoothingFunction().method4 + bleu_100 = corpus_bleu(references, hypotheses, smoothing_function=sf) * 100.0 + + print(f" Corpus BLEU-4: {bleu_100:.2f} | elapsed: {elapsed:.1f}s") + return bleu_100, elapsed + + # ── Unified evaluate entry point ───────────────────────────────────────── + + def evaluate( + self, + dataset, + pool_size: int = POOL_SIZE, + max_samples: Optional[int] = None, + ) -> dict: + """Evaluate and return a result dict. + + Dispatches to :meth:`compute_mrr` or :meth:`compute_bleu` depending + on ``self.task``. + + Returns: + Dict with keys: + - ``'task'`` : str + - ``'model_type'`` : str + - ``'simplification_ratio'``: float + - ``'mrr'`` : float (search only) + - ``'bleu4'`` : float (summarization only) + - ``'elapsed'`` : float + """ + if self.task == "search": + metric, elapsed = self.compute_mrr( + dataset, pool_size=pool_size, max_samples=max_samples + ) + return { + "task": self.task, + "model_type": self.model_type, + "simplification_ratio": self.simplification_ratio, + "mrr": metric, + "elapsed": elapsed, + } + metric, elapsed = self.compute_bleu(dataset, max_samples=max_samples) + return { + "task": self.task, + "model_type": self.model_type, + "simplification_ratio": self.simplification_ratio, + "bleu4": metric, + "elapsed": elapsed, + } + + +# ───────────────────────────────────────────────────────────────────────────── +# Full evaluation sweep across multiple simplification ratios +# ───────────────────────────────────────────────────────────────────────────── + +def evaluate_dynamic_sweep( + task: str, + model_type: str, + dynamic_checkpoint: str, + downstream_checkpoint: Optional[str], + ratios: list[float] = RATIOS, + language: str = "java", + split: str = "test", + max_samples: Optional[int] = None, + pool_size: int = POOL_SIZE, + device: Optional[str] = None, + output_file: Optional[str] = None, +) -> list[dict]: + """Run evaluation for multiple simplification ratios and print a table. + + Args: + task: ``'search'`` or ``'summarization'``. + model_type: ``'codebert'`` or ``'codet5'``. + dynamic_checkpoint: Path to the dynamic model ``.pt`` checkpoint. + downstream_checkpoint: Path to the fine-tuned downstream checkpoint. + ratios: Simplification ratios to evaluate. + language: CodeSearchNet language. + split: Dataset split. + max_samples: Limit evaluation to first N samples. + pool_size: Pool size for MRR (search only). + device: Device string. + output_file: Optional JSON file to write results to. + + Returns: + List of result dicts, one per ratio. + """ + dataset = load_codesearchnet(language)[split] + rows: list[dict] = [] + + metric_key = "mrr" if task == "search" else "bleu4" + + for ratio in ratios: + print(f"\n{'='*60}\nRatio {int(ratio*100)}%\n{'='*60}") + evaluator = DynamicEvaluator( + task = task, + model_type = model_type, + dynamic_checkpoint = dynamic_checkpoint, + downstream_checkpoint = downstream_checkpoint, + simplification_ratio = ratio, + device = device, + ) + result = evaluator.evaluate(dataset, pool_size=pool_size, max_samples=max_samples) + rows.append(result) + print(f" {metric_key.upper()}={result[metric_key]:.4f} elapsed={result['elapsed']:.1f}s") + + # Print summary table + metric_name = "MRR" if task == "search" else "BLEU-4" + metric_fmt = ".4f" if task == "search" else ".2f" + print(f"\n{'='*50}") + print(f"Dynamic {task.upper()} ({model_type}) — {metric_name}") + print(f"{'Ratio':<8} {metric_name:>10} {'Time (s)':>10}") + print("-" * 32) + for row in rows: + ratio_str = f"{int(row['simplification_ratio']*100)}%" + print( + f"{ratio_str:<8} " + f"{format(row[metric_key], metric_fmt):>10} " + f"{row['elapsed']:>10.1f}" + ) + print("=" * 50) + + if output_file: + Path(output_file).parent.mkdir(parents=True, exist_ok=True) + with open(output_file, "w", encoding="utf-8") as f: + json.dump({"task": task, "model": model_type, "rows": rows}, f, indent=2) + print(f"Results saved → {output_file}") + + return rows + + +# ───────────────────────────────────────────────────────────────────────────── +# CLI +# ───────────────────────────────────────────────────────────────────────────── + +def _build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + description="Evaluate integrated dynamic model on MRR or BLEU-4", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument("--task", choices=["search", "summarization"], default="search") + p.add_argument("--model", choices=["codebert", "codet5"], default="codebert") + p.add_argument("--dynamic-checkpoint", required=True, + help="Path to the dynamic model .pt checkpoint.") + p.add_argument("--downstream-checkpoint", default=None, + help="Fine-tuned CodeBERT/CodeT5 checkpoint directory.") + p.add_argument("--simplification-ratio", type=float, default=0.30) + p.add_argument("--ratios", nargs="+", type=float, default=None, + help="Multiple ratios for a sweep (overrides --simplification-ratio).") + p.add_argument("--language", default="java") + p.add_argument("--split", default="test", choices=["train", "validation", "test"]) + p.add_argument("--max-samples", type=int, default=None) + p.add_argument("--pool-size", type=int, default=POOL_SIZE) + p.add_argument("--device", default=None) + p.add_argument("--output", default=None, metavar="FILE") + return p + + +def main() -> None: + """CLI entry point.""" + args = _build_arg_parser().parse_args() + + if args.ratios: + evaluate_dynamic_sweep( + task = args.task, + model_type = args.model, + dynamic_checkpoint = args.dynamic_checkpoint, + downstream_checkpoint = args.downstream_checkpoint, + ratios = args.ratios, + language = args.language, + split = args.split, + max_samples = args.max_samples, + pool_size = args.pool_size, + device = args.device, + output_file = args.output, + ) + else: + dataset = load_codesearchnet(args.language)[args.split] + evaluator = DynamicEvaluator( + task = args.task, + model_type = args.model, + dynamic_checkpoint = args.dynamic_checkpoint, + downstream_checkpoint = args.downstream_checkpoint, + simplification_ratio = args.simplification_ratio, + device = args.device, + ) + result = evaluator.evaluate(dataset, pool_size=args.pool_size, max_samples=args.max_samples) + print(f"\nResult: {result}") + + if args.output: + Path(args.output).parent.mkdir(parents=True, exist_ok=True) + with open(args.output, "w", encoding="utf-8") as f: + json.dump(result, f, indent=2) + print(f"Results saved → {args.output}") + + +if __name__ == "__main__": + main() diff --git a/src/models/dynamic/__init__.py b/src/models/dynamic/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/dynamic/cnn_model.py b/src/models/dynamic/cnn_model.py new file mode 100644 index 0000000..3741488 --- /dev/null +++ b/src/models/dynamic/cnn_model.py @@ -0,0 +1,104 @@ +""" +CNN Dynamic Importance Scoring Model +====================================== +A 1-D convolutional network that simultaneously scores all tokens in a code +snippet. The configurable ``kernel_size`` controls how many neighbouring +tokens each position can "see" before predicting its importance score, +providing local context without the sequential dependency of an RNN. + +Architecture +------------ +Given a token sequence of length *T*: + + 1. Token embedding: E ∈ ℝ^{T × embed_dim} + 2. 1-D convolution stack (each layer uses same-padding to preserve T): + Conv1d(embed_dim, hidden_dim, kernel_size) → ReLU + Conv1d(hidden_dim, hidden_dim, kernel_size) → ReLU + Conv1d(hidden_dim, 1, 1) → Sigmoid + 3. Output: importance scores s ∈ (0, 1)^T + +Trained with MSE loss against per-token attention scores from a fine-tuned +transformer (produced by ``src/data/dynamic/dynamic_dataset_generator.py``). + +Usage +----- + from src.models.dynamic.cnn_model import CNNScoringModel + + model = CNNScoringModel(kernel_size=5) + + # token_ids: LongTensor [batch, seq_len] + scores = model(token_ids) # FloatTensor [batch, seq_len] +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + + +class CNNScoringModel(nn.Module): + """1-D CNN importance scorer with configurable receptive field. + + Args: + vocab_size: Tokenizer vocabulary size. + embed_dim: Dimension of the token embedding. + hidden_dim: Number of feature maps in each convolutional layer. + num_layers: Number of Conv1d → ReLU layers before the output head. + kernel_size: Convolutional kernel size; must be odd so same-padding + is exact. Larger values increase the receptive field. + dropout: Dropout probability applied after each hidden conv layer. + + Usage:: + + model = CNNScoringModel(kernel_size=5) + scores = model(token_ids) # [batch, seq_len] + """ + + def __init__( + self, + vocab_size: int = 50265, + embed_dim: int = 64, + hidden_dim: int = 128, + num_layers: int = 2, + kernel_size: int = 5, + dropout: float = 0.1, + ) -> None: + super().__init__() + + if kernel_size % 2 == 0: + raise ValueError("kernel_size must be odd to ensure same-length output.") + + self.embedding = nn.Embedding(vocab_size, embed_dim) + + padding = kernel_size // 2 + conv_layers: list[nn.Module] = [] + in_channels = embed_dim + for _ in range(num_layers): + conv_layers += [ + nn.Conv1d(in_channels, hidden_dim, kernel_size, padding=padding), + nn.ReLU(), + nn.Dropout(dropout), + ] + in_channels = hidden_dim + # Output head: 1×1 conv collapses channels to a single score per position + conv_layers += [ + nn.Conv1d(in_channels, 1, kernel_size=1), + nn.Sigmoid(), + ] + + self.conv_stack = nn.Sequential(*conv_layers) + + def forward(self, token_ids: torch.Tensor) -> torch.Tensor: + """Compute importance scores for a batch of token sequences. + + Args: + token_ids: LongTensor of token IDs with shape ``[batch, seq_len]``. + + Returns: + FloatTensor of per-token scores in ``(0, 1)``, + shape ``[batch, seq_len]``. + """ + # [B, T, embed_dim] → transpose to [B, embed_dim, T] for Conv1d + x = self.embedding(token_ids).transpose(1, 2) + out = self.conv_stack(x) # [B, 1, T] + return out.squeeze(1) # [B, T] diff --git a/src/models/dynamic/linear_model.py b/src/models/dynamic/linear_model.py new file mode 100644 index 0000000..64dbf67 --- /dev/null +++ b/src/models/dynamic/linear_model.py @@ -0,0 +1,121 @@ +""" +Linear Dynamic Importance Scoring Model +======================================== +A multi-hidden-layer MLP that assigns an importance score to each individual +token based on its identity, position in the sequence, and enclosing AST +statement category. The model processes each token *independently* — there +is no cross-token context, making it the lightest of the three dynamic +architectures proposed in the dynamic approaches plan. + +Architecture +------------ +For each token t in a code snippet: + + 1. Look up a learned token embedding: e_tok ∈ ℝ^{token_embed_dim} + 2. Look up a learned category embedding: e_cat ∈ ℝ^{cat_embed_dim} + 3. Compute a scalar position feature: p = position / max_seq_len + + 4. Concatenate: x = [e_tok ‖ e_cat ‖ p] ∈ ℝ^{token_embed_dim + cat_embed_dim + 1} + 5. MLP: + fc1(x) → ReLU → Dropout + fc2 → ReLU → Dropout + fc3 → Sigmoid + → importance score s ∈ (0, 1) + +Trained with MSE loss against per-token attention scores from a fine-tuned +transformer (produced by ``src/data/dynamic/dynamic_dataset_generator.py``). + +Usage +----- + from src.models.dynamic.linear_model import LinearScoringModel + + model = LinearScoringModel() + + # token_ids, positions, category_ids: LongTensors [N] (N = seq_len) + scores = model(token_ids, positions, category_ids) # FloatTensor [N] + + # Batched inference (pad to same length first): + scores = model(token_ids_2d, positions_2d, category_ids_2d) # [B, seq_len] +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn as nn + + +class LinearScoringModel(nn.Module): + """Per-token MLP importance scorer. + + Args: + vocab_size: Tokenizer vocabulary size. + num_categories: Number of AST statement categories. + token_embed_dim: Dimension of the token embedding. + cat_embed_dim: Dimension of the category embedding. + hidden_dim: Width of each hidden fully-connected layer. + num_hidden: Number of hidden layers (minimum 1). + dropout: Dropout probability applied after each hidden layer. + max_seq_len: Used to normalise position to [0, 1]. + + Usage:: + + model = LinearScoringModel(vocab_size=50265, num_categories=28) + scores = model(token_ids, positions, category_ids) # [N] or [B, N] + """ + + def __init__( + self, + vocab_size: int = 50265, + num_categories: int = 28, + token_embed_dim: int = 64, + cat_embed_dim: int = 16, + hidden_dim: int = 128, + num_hidden: int = 2, + dropout: float = 0.1, + max_seq_len: int = 512, + ) -> None: + super().__init__() + self.max_seq_len = max_seq_len + + self.token_embedding = nn.Embedding(vocab_size, token_embed_dim) + self.category_embedding = nn.Embedding(num_categories, cat_embed_dim) + + input_dim = token_embed_dim + cat_embed_dim + 1 # +1 for position scalar + + layers: list[nn.Module] = [] + in_dim = input_dim + for _ in range(num_hidden): + layers += [nn.Linear(in_dim, hidden_dim), nn.ReLU(), nn.Dropout(dropout)] + in_dim = hidden_dim + layers += [nn.Linear(in_dim, 1), nn.Sigmoid()] + + self.mlp = nn.Sequential(*layers) + + def forward( + self, + token_ids: torch.Tensor, + positions: torch.Tensor, + category_ids: torch.Tensor, + ) -> torch.Tensor: + """Compute importance scores for a (batch of) token sequence(s). + + Args: + token_ids: LongTensor of token IDs, shape ``[N]`` or ``[B, N]``. + positions: LongTensor of position indices, same shape. + category_ids: LongTensor of category IDs, same shape. + + Returns: + FloatTensor of scores in ``(0, 1)``, shape ``[N]`` or ``[B, N]``. + """ + e_tok = self.token_embedding(token_ids) # [..., token_embed_dim] + e_cat = self.category_embedding(category_ids) # [..., cat_embed_dim] + pos = positions.float() / self.max_seq_len # [...], normalised + + # Append position as a trailing scalar dimension + pos_feat = pos.unsqueeze(-1) # [..., 1] + x = torch.cat([e_tok, e_cat, pos_feat], dim=-1) # [..., input_dim] + + out = self.mlp(x) # [..., 1] + return out.squeeze(-1) # [...] — shape matches input diff --git a/src/models/dynamic/rnn_model.py b/src/models/dynamic/rnn_model.py new file mode 100644 index 0000000..7351ae9 --- /dev/null +++ b/src/models/dynamic/rnn_model.py @@ -0,0 +1,121 @@ +""" +RNN Dynamic Importance Scoring Model +====================================== +A bidirectional GRU that scores all tokens in a code snippet simultaneously. +The bidirectional architecture allows each token to incorporate context from +both directions (left and right) when predicting its importance score, +providing full-sequence context at a fraction of the cost of a transformer. + +Architecture +------------ +Given a token sequence of length *T*: + + 1. Token embedding: E ∈ ℝ^{T × embed_dim} + 2. Bidirectional GRU: + BiGRU(embed_dim → hidden_dim×2) with ``num_layers`` layers + 3. Output projection: + Linear(hidden_dim × 2, 1) → Sigmoid + 4. Output: importance scores s ∈ (0, 1)^T + +Trained with MSE loss against per-token attention scores from a fine-tuned +transformer (produced by ``src/data/dynamic/dynamic_dataset_generator.py``). + +Usage +----- + from src.models.dynamic.rnn_model import RNNScoringModel + + model = RNNScoringModel() + + # token_ids: LongTensor [batch, seq_len] + scores = model(token_ids) # FloatTensor [batch, seq_len] + + # With sequence lengths for packed-sequence efficiency: + scores = model(token_ids, lengths=lengths) # FloatTensor [batch, seq_len] +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn as nn +from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence + + +class RNNScoringModel(nn.Module): + """Bidirectional GRU importance scorer. + + Args: + vocab_size: Tokenizer vocabulary size. + embed_dim: Dimension of the token embedding. + hidden_dim: GRU hidden state size *per direction* (output is + ``hidden_dim × 2`` due to bidirectionality). + num_layers: Number of stacked GRU layers. + dropout: Dropout between GRU layers (applied when + ``num_layers > 1``). + + Usage:: + + model = RNNScoringModel() + scores = model(token_ids) # [batch, seq_len] + scores = model(token_ids, lengths) # packed-sequence path + """ + + def __init__( + self, + vocab_size: int = 50265, + embed_dim: int = 64, + hidden_dim: int = 128, + num_layers: int = 2, + dropout: float = 0.1, + ) -> None: + super().__init__() + + self.embedding = nn.Embedding(vocab_size, embed_dim) + self.gru = nn.GRU( + input_size = embed_dim, + hidden_size = hidden_dim, + num_layers = num_layers, + batch_first = True, + bidirectional = True, + dropout = dropout if num_layers > 1 else 0.0, + ) + # Project bidirectional hidden state to a single score + self.output_head = nn.Sequential( + nn.Linear(hidden_dim * 2, 1), + nn.Sigmoid(), + ) + + def forward( + self, + token_ids: torch.Tensor, + lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Compute importance scores for a batch of token sequences. + + Args: + token_ids: LongTensor with shape ``[batch, seq_len]``. + lengths: Optional LongTensor ``[batch]`` with the actual + (non-padded) length of each sequence. When supplied, + packed sequences are used for efficiency and scores + at padding positions are set to zero. + + Returns: + FloatTensor of per-token scores in ``(0, 1)``, + shape ``[batch, seq_len]``. + """ + x = self.embedding(token_ids) # [B, T, embed_dim] + + if lengths is not None: + # Use packed sequences to avoid computing over padding tokens + lengths_cpu = lengths.cpu() + packed = pack_padded_sequence( + x, lengths_cpu, batch_first=True, enforce_sorted=False + ) + packed_out, _ = self.gru(packed) + output, _ = pad_packed_sequence(packed_out, batch_first=True) + else: + output, _ = self.gru(x) # [B, T, hidden_dim * 2] + + scores = self.output_head(output) # [B, T, 1] + return scores.squeeze(-1) # [B, T] diff --git a/src/models/integrated/__init__.py b/src/models/integrated/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/integrated/integrated_codebert.py b/src/models/integrated/integrated_codebert.py new file mode 100644 index 0000000..ed422f5 --- /dev/null +++ b/src/models/integrated/integrated_codebert.py @@ -0,0 +1,274 @@ +""" +Integrated Dynamic CodeBERT +============================= +Wraps ``FinetunedCodeBERT`` with a dynamic token-importance scoring step +so that code simplification happens *at inference time* rather than from a +pre-computed static score dictionary. + +Pipeline (per forward call) +--------------------------- +1. Tokenise the raw code snippet. +2. Run the lightweight dynamic scoring model (Linear / CNN / RNN) to + obtain a per-token importance score. +3. Remove the ``simplification_ratio`` fraction of lowest-scored tokens by + masking their character spans in the source string. +4. Re-tokenise the simplified code and pass it to ``FinetunedCodeBERT``. + +The dynamic model must have been trained via +``src/training/dynamic/train_dynamic.py`` and saved as a ``.pt`` checkpoint. + +Usage +----- + from src.models.integrated.integrated_codebert import IntegratedCodeBERT + + model = IntegratedCodeBERT( + dynamic_checkpoint = "checkpoints/dynamic/rnn_codebert_search/best.pt", + codebert_checkpoint = "checkpoints/codebert_encoder_only/epoch_1", + simplification_ratio = 0.30, + ) + + # Classification (code search) + logits = model.forward_classification(input_ids, attention_mask) + + # Generation (code summarization) + summary = model.forward_generation("public int add(int a, int b) { return a + b; }") +""" + +from __future__ import annotations + +import re +from typing import Optional, cast + +import torch +from transformers.models.roberta import RobertaTokenizer + +from src.models.finetuned_codebert_model import FinetunedCodeBERT +from src.training.dynamic.train_dynamic import load_dynamic_model +from src.data.dynamic.dynamic_dataset_generator import ( + ALL_CATEGORIES, + CATEGORY_TO_ID, +) +from src.attention.attention_extraction import ( + map_subword_tokens_to_categories, + SPECIAL_TOKEN_STRINGS, +) + +MAX_LENGTH = 512 +SPECIAL_TOKENS = {"[CLS]", "[SEP]", "", "", "", ""} + + +def _simplify_code( + code: str, + tokenizer: RobertaTokenizer, + dynamic_model: torch.nn.Module, + arch: str, + simplification_ratio: float, + device: torch.device, +) -> str: + """Use the dynamic model to score tokens and return simplified code. + + Args: + code: Raw source code string. + tokenizer: Tokenizer for the downstream model. + dynamic_model: Pre-loaded dynamic scoring model. + arch: Architecture tag (``'linear'``, ``'cnn'``, ``'rnn'``). + simplification_ratio: Fraction of removable tokens to drop. + device: Device for dynamic model inference. + + Returns: + Simplified code string with low-importance tokens removed. + """ + if not code.strip(): + return code + + enc = tokenizer( + code, + truncation=True, + max_length=MAX_LENGTH, + return_offsets_mapping=True, + add_special_tokens=True, + ) + token_ids_list = enc["input_ids"] + offset_mapping = enc["offset_mapping"] + tokens = tokenizer.convert_ids_to_tokens(token_ids_list) + categories = map_subword_tokens_to_categories(code, tokens) + category_ids = [ + CATEGORY_TO_ID.get(cat, CATEGORY_TO_ID["Other"]) for cat in categories + ] + seq_len = len(token_ids_list) + positions = list(range(seq_len)) + + token_ids_t = torch.tensor(token_ids_list, dtype=torch.long, device=device).unsqueeze(0) + positions_t = torch.tensor(positions, dtype=torch.long, device=device).unsqueeze(0) + category_ids_t = torch.tensor(category_ids, dtype=torch.long, device=device).unsqueeze(0) + lengths_t = torch.tensor([seq_len], dtype=torch.long, device=device) + + with torch.no_grad(): + if arch == "linear": + scores = dynamic_model( + token_ids_t.squeeze(0), + positions_t.squeeze(0), + category_ids_t.squeeze(0), + ).tolist() + elif arch == "rnn": + scores = dynamic_model(token_ids_t, lengths_t).squeeze(0).tolist() + else: # cnn + scores = dynamic_model(token_ids_t).squeeze(0).tolist() + + # Identify removable positions (non-special tokens only) + removable_indices = [ + i for i, tok in enumerate(tokens) + if tok not in SPECIAL_TOKENS + ] + n_remove = int(simplification_ratio * len(removable_indices)) + + if n_remove == 0: + return code + + # Sort removable positions by ascending score, take lowest n_remove + removable_sorted = sorted(removable_indices, key=lambda i: scores[i]) + removed_set = set(removable_sorted[:n_remove]) + + # Mask character spans for removed tokens + removed_chars: set[int] = set() + for idx in removed_set: + if idx >= len(offset_mapping): + continue + char_start, char_end = offset_mapping[idx] + if char_start == char_end: # special token — no source span + continue + removed_chars.update(range(char_start, char_end)) + + simplified = "".join(ch for i, ch in enumerate(code) if i not in removed_chars) + simplified = re.sub(r"\n{3,}", "\n\n", simplified) + return simplified + + +class IntegratedCodeBERT(FinetunedCodeBERT): + """FinetunedCodeBERT with dynamic token-importance simplification. + + Inherits all methods from ``FinetunedCodeBERT`` and overrides + ``forward_classification`` and ``forward_generation`` to prepend the + dynamic simplification step. + + Args: + dynamic_checkpoint: Path to a ``.pt`` checkpoint produced by + ``train_dynamic.py``. + codebert_checkpoint: Path to a fine-tuned CodeBERT checkpoint (as + produced by ``FinetunedCodeBERT.save_model()``). + If ``None``, pretrained weights are used. + simplification_ratio: Fraction of removable tokens to drop + (e.g. ``0.30`` removes the 30% lowest-scored). + device: Device string; auto-detected if ``None``. + + Usage:: + + model = IntegratedCodeBERT( + dynamic_checkpoint = "checkpoints/dynamic/rnn_codebert_search/best.pt", + codebert_checkpoint = "checkpoints/codebert_encoder_only/epoch_1", + simplification_ratio = 0.30, + ) + logits = model.forward_classification(input_ids, attention_mask) + """ + + def __init__( + self, + dynamic_checkpoint: str, + codebert_checkpoint: Optional[str] = None, + simplification_ratio: float = 0.30, + device: Optional[str] = None, + ) -> None: + super().__init__(device=device) + + if codebert_checkpoint is not None: + self.load_checkpoint(codebert_checkpoint) + + self.simplification_ratio = simplification_ratio + + # Load the dynamic scoring model (load_dynamic_model handles device placement) + self._dynamic_model = load_dynamic_model( + dynamic_checkpoint, device_str=device or self.device + ) + + # Determine architecture tag from the checkpoint metadata + ckpt = torch.load(dynamic_checkpoint, map_location=self.device) + self._dynamic_arch: str = ckpt.get("arch", "rnn") + + # ------------------------------------------------------------------ + # Simplification helper + # ------------------------------------------------------------------ + + def _simplify(self, code: str) -> str: + """Return dynamically simplified code.""" + _, tokenizer = self._load_model() + return _simplify_code( + code = code, + tokenizer = cast(RobertaTokenizer, tokenizer), + dynamic_model = self._dynamic_model, + arch = self._dynamic_arch, + simplification_ratio = self.simplification_ratio, + device = torch.device(self.device), + ) + + # ------------------------------------------------------------------ + # Overridden forward methods + # ------------------------------------------------------------------ + + def forward_classification( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + classification_labels: Optional[torch.Tensor] = None, + ): + """Classify (query, simplified_code) pairs. + + The classification head operates on the tokenized pair as passed in + ``input_ids``; simplification is applied during the *generation* path + where raw code strings are available. For classification, the caller + is expected to pass pre-simplified input when desired, or call + ``simplify_and_classify`` for the end-to-end path. + + Delegates directly to ``FinetunedCodeBERT.forward_classification``. + """ + return super().forward_classification( + input_ids = input_ids, + attention_mask = attention_mask, + classification_labels = classification_labels, + ) + + def forward_generation( + self, + code: str, + max_src_len: int = 512, + max_new_tokens: int = 64, + num_beams: int = 1, + ) -> str: + """Dynamically simplify then generate a summary. + + Args: + code: Raw source code string. + max_src_len: Max encoder input length. + max_new_tokens: Max decoder output tokens. + num_beams: Beam width (only greedy decoding is supported). + + Returns: + Decoded summary string produced from simplified code. + """ + simplified = self._simplify(code) + return super().forward_generation( + code = simplified if simplified.strip() else code, + max_src_len = max_src_len, + max_new_tokens = max_new_tokens, + num_beams = num_beams, + ) + + def simplify_code(self, code: str) -> str: + """Return the dynamically simplified version of a code snippet. + + Args: + code: Raw source code string. + + Returns: + Simplified code string. + """ + return self._simplify(code) diff --git a/src/models/integrated/integrated_codet5.py b/src/models/integrated/integrated_codet5.py new file mode 100644 index 0000000..cf64c33 --- /dev/null +++ b/src/models/integrated/integrated_codet5.py @@ -0,0 +1,249 @@ +""" +Integrated Dynamic CodeT5 +=========================== +Wraps ``FineTunedCodeT5`` with a dynamic token-importance scoring step so +that code simplification happens *at inference time* rather than from a +pre-computed static score dictionary. + +Pipeline (per forward call) +--------------------------- +1. Tokenise the raw code snippet. +2. Run the lightweight dynamic scoring model (Linear / CNN / RNN) to + obtain a per-token importance score. +3. Remove the ``simplification_ratio`` fraction of lowest-scored tokens by + masking their character spans in the source string. +4. Re-tokenise the simplified code and pass it to ``FineTunedCodeT5``. + +The dynamic model must have been trained via +``src/training/dynamic/train_dynamic.py`` and saved as a ``.pt`` checkpoint. + +Usage +----- + from src.models.integrated.integrated_codet5 import IntegratedCodeT5 + + model = IntegratedCodeT5( + dynamic_checkpoint = "checkpoints/dynamic/rnn_codet5_summarization/best.pt", + codet5_checkpoint = "checkpoints/codet5_encoder_decoder/epoch_1", + simplification_ratio = 0.30, + ) + + # Summarisation (code → natural language) + summary = model.forward_generation("public int add(int a, int b) { return a + b; }") + + # Classification (code search) + logits = model.forward_classification(input_ids, attention_mask) +""" + +from __future__ import annotations + +import re +from typing import Any, Optional, cast + +import torch +from transformers.models.roberta import RobertaTokenizer + +from src.models.finetuned_codet5_model import FineTunedCodeT5 +from src.training.dynamic.train_dynamic import load_dynamic_model +from src.data.dynamic.dynamic_dataset_generator import CATEGORY_TO_ID +from src.attention.attention_extraction import ( + map_subword_tokens_to_categories, + SPECIAL_TOKEN_STRINGS, +) + +MAX_LENGTH = 512 +SPECIAL_TOKENS = {"[CLS]", "[SEP]", "", "", "", ""} + + +def _simplify_code( + code: str, + tokenizer: Any, + dynamic_model: torch.nn.Module, + arch: str, + simplification_ratio: float, + device: torch.device, +) -> str: + """Use the dynamic model to score tokens and return simplified code. + + Args: + code: Raw source code string. + tokenizer: Tokenizer for the downstream model. + dynamic_model: Pre-loaded dynamic scoring model. + arch: Architecture tag (``'linear'``, ``'cnn'``, ``'rnn'``). + simplification_ratio: Fraction of removable tokens to drop. + device: Device for dynamic model inference. + + Returns: + Simplified code string with low-importance tokens removed. + """ + if not code.strip(): + return code + + enc = tokenizer( + code, + truncation=True, + max_length=MAX_LENGTH, + return_offsets_mapping=True, + add_special_tokens=True, + ) + token_ids_list = enc["input_ids"] + offset_mapping = enc["offset_mapping"] + tokens = tokenizer.convert_ids_to_tokens(token_ids_list) + categories = map_subword_tokens_to_categories(code, tokens) + category_ids = [ + CATEGORY_TO_ID.get(cat, CATEGORY_TO_ID["Other"]) for cat in categories + ] + seq_len = len(token_ids_list) + positions = list(range(seq_len)) + + token_ids_t = torch.tensor(token_ids_list, dtype=torch.long, device=device).unsqueeze(0) + positions_t = torch.tensor(positions, dtype=torch.long, device=device).unsqueeze(0) + category_ids_t = torch.tensor(category_ids, dtype=torch.long, device=device).unsqueeze(0) + lengths_t = torch.tensor([seq_len], dtype=torch.long, device=device) + + with torch.no_grad(): + if arch == "linear": + scores = dynamic_model( + token_ids_t.squeeze(0), + positions_t.squeeze(0), + category_ids_t.squeeze(0), + ).tolist() + elif arch == "rnn": + scores = dynamic_model(token_ids_t, lengths_t).squeeze(0).tolist() + else: # cnn + scores = dynamic_model(token_ids_t).squeeze(0).tolist() + + # Identify removable positions (non-special tokens only) + removable_indices = [ + i for i, tok in enumerate(tokens) + if tok not in SPECIAL_TOKENS + ] + n_remove = int(simplification_ratio * len(removable_indices)) + + if n_remove == 0: + return code + + # Sort removable positions by ascending score, take lowest n_remove + removable_sorted = sorted(removable_indices, key=lambda i: scores[i]) + removed_set = set(removable_sorted[:n_remove]) + + # Mask character spans for removed tokens + removed_chars: set[int] = set() + for idx in removed_set: + if idx >= len(offset_mapping): + continue + char_start, char_end = offset_mapping[idx] + if char_start == char_end: # special token — no source span + continue + removed_chars.update(range(char_start, char_end)) + + simplified = "".join(ch for i, ch in enumerate(code) if i not in removed_chars) + simplified = re.sub(r"\n{3,}", "\n\n", simplified) + return simplified + + +class IntegratedCodeT5(FineTunedCodeT5): + """FineTunedCodeT5 with dynamic token-importance simplification. + + Inherits all methods from ``FineTunedCodeT5`` and overrides + ``forward_generation`` to prepend the dynamic simplification step. + ``forward_classification`` is also available from the parent. + + Args: + dynamic_checkpoint: Path to a ``.pt`` checkpoint produced by + ``train_dynamic.py``. + codet5_checkpoint: Path to a fine-tuned CodeT5 checkpoint (as + produced by ``FineTunedCodeT5.save_model()``). + If ``None``, pretrained weights are used. + simplification_ratio: Fraction of removable tokens to drop + (e.g. ``0.30`` removes the 30% lowest-scored). + device: Device string; auto-detected if ``None``. + + Usage:: + + model = IntegratedCodeT5( + dynamic_checkpoint = "checkpoints/dynamic/rnn_codet5_summ/best.pt", + codet5_checkpoint = "checkpoints/codet5_encoder_decoder/epoch_1", + simplification_ratio = 0.30, + ) + summary = model.forward_generation("public int add(int a, int b) { ... }") + """ + + def __init__( + self, + dynamic_checkpoint: str, + codet5_checkpoint: Optional[str] = None, + simplification_ratio: float = 0.30, + device: Optional[str] = None, + ) -> None: + super().__init__(device=device) + + if codet5_checkpoint is not None: + self.load_checkpoint(codet5_checkpoint) + + self.simplification_ratio = simplification_ratio + + # Load the dynamic scoring model (load_dynamic_model handles device placement) + self._dynamic_model = load_dynamic_model( + dynamic_checkpoint, device_str=device or self.device + ) + + # Determine architecture tag from the checkpoint metadata + ckpt = torch.load(dynamic_checkpoint, map_location=self.device) + self._dynamic_arch: str = ckpt.get("arch", "rnn") + + # ------------------------------------------------------------------ + # Simplification helper + # ------------------------------------------------------------------ + + def _simplify(self, code: str) -> str: + """Return dynamically simplified code.""" + _, tokenizer = self._load_model() + return _simplify_code( + code = code, + tokenizer = tokenizer, + dynamic_model = self._dynamic_model, + arch = self._dynamic_arch, + simplification_ratio = self.simplification_ratio, + device = torch.device(self.device), + ) + + # ------------------------------------------------------------------ + # Overridden forward methods + # ------------------------------------------------------------------ + + def forward_generation( + self, + code: str, + max_src_len: int = 512, + max_new_tokens: int = 64, + num_beams: int = 4, + ) -> str: + """Dynamically simplify then generate a natural-language summary. + + Args: + code: Raw source code string. + max_src_len: Max encoder input length. + max_new_tokens: Max decoder output tokens. + num_beams: Beam search width (4 matches the paper setting). + + Returns: + Decoded summary string produced from the simplified code. + """ + simplified = self._simplify(code) + return super().forward_generation( + code = simplified if simplified.strip() else code, + max_src_len = max_src_len, + max_new_tokens = max_new_tokens, + num_beams = num_beams, + ) + + def simplify_code(self, code: str) -> str: + """Return the dynamically simplified version of a code snippet. + + Args: + code: Raw source code string. + + Returns: + Simplified code string. + """ + return self._simplify(code) diff --git a/src/training/dynamic/__init__.py b/src/training/dynamic/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/training/dynamic/train_dynamic.py b/src/training/dynamic/train_dynamic.py new file mode 100644 index 0000000..3c87542 --- /dev/null +++ b/src/training/dynamic/train_dynamic.py @@ -0,0 +1,452 @@ +""" +Dynamic Model Training +======================= +Trains one of the three lightweight dynamic importance-scoring models +(Linear, CNN, or RNN) on the dataset generated by +``src/data/dynamic/dynamic_dataset_generator.py``. + +All three models are trained to minimise MSE between their predicted +per-token scores and the attention scores extracted from a fine-tuned +transformer (CodeBERT or CodeT5). After each epoch the best checkpoint +(lowest validation MSE) is saved. + +Architecture selection +---------------------- ++----------+---------------------------------------+ +| arch | model class | ++==========+=======================================+ +| linear | ``LinearScoringModel`` | ++----------+---------------------------------------+ +| cnn | ``CNNScoringModel`` | ++----------+---------------------------------------+ +| rnn | ``RNNScoringModel`` | ++----------+---------------------------------------+ + +Usage +----- +CLI (from repository root):: + + python -m src.training.dynamic.train_dynamic \\ + --arch rnn \\ + --train-data data/dynamic/dynamic_codebert_search.json \\ + --val-data data/dynamic/dynamic_codebert_search_val.json \\ + --save-dir checkpoints/dynamic/rnn_codebert_search \\ + --epochs 10 --batch-size 64 + +Programmatic:: + + from src.training.dynamic.train_dynamic import train_dynamic + + ckpt_path = train_dynamic( + arch = "rnn", + train_data = "data/dynamic/dynamic_codebert_search.json", + epochs = 5, + ) +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import Optional + +import torch +import torch.nn as nn +from torch.optim import AdamW +from torch.utils.data import DataLoader, random_split +from transformers.optimization import get_linear_schedule_with_warmup + +from src.data.dynamic.dynamic_dataset_generator import ( + DynamicScoringDataset, + dynamic_collate_fn, +) +from src.models.dynamic.linear_model import LinearScoringModel +from src.models.dynamic.cnn_model import CNNScoringModel +from src.models.dynamic.rnn_model import RNNScoringModel + +_ARCH_MAP = { + "linear": LinearScoringModel, + "cnn": CNNScoringModel, + "rnn": RNNScoringModel, +} + + +# --------------------------------------------------------------------------- +# Loss helper +# --------------------------------------------------------------------------- + +def _masked_mse( + predictions: torch.Tensor, + targets: torch.Tensor, + padding_mask: torch.Tensor, +) -> torch.Tensor: + """MSE loss over non-padding positions only. + + Args: + predictions: FloatTensor ``[B, T]``. + targets: FloatTensor ``[B, T]``. + padding_mask: BoolTensor ``[B, T]``, True where the token is real. + + Returns: + Scalar MSE loss averaged over all valid (non-padded) positions. + """ + diff = (predictions - targets) ** 2 + masked = diff[padding_mask] + return masked.mean() + + +# --------------------------------------------------------------------------- +# Single-epoch training / validation helpers +# --------------------------------------------------------------------------- + +def _run_epoch( + model: nn.Module, + loader: DataLoader, + optimizer: Optional[torch.optim.Optimizer], + scheduler, + arch: str, + device: torch.device, + train: bool, + max_grad_norm: float = 1.0, +) -> float: + """Run one training or validation epoch and return the average MSE.""" + model.train(train) + total_loss = 0.0 + n_batches = 0 + + ctx = torch.enable_grad() if train else torch.no_grad() + with ctx: + for batch in loader: + token_ids = batch["token_ids"].to(device) + positions = batch["positions"].to(device) + category_ids = batch["category_ids"].to(device) + attention_scores = batch["attention_scores"].to(device) + padding_mask = batch["padding_mask"].to(device) + lengths = batch["lengths"].to(device) + + if arch == "linear": + preds = model(token_ids, positions, category_ids) + elif arch == "cnn": + preds = model(token_ids) + else: # rnn + preds = model(token_ids, lengths) + + loss = _masked_mse(preds, attention_scores, padding_mask) + + if train and optimizer is not None: + optimizer.zero_grad() + loss.backward() + nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) + optimizer.step() + if scheduler is not None: + scheduler.step() + + total_loss += loss.item() + n_batches += 1 + + return total_loss / max(n_batches, 1) + + +# --------------------------------------------------------------------------- +# Public training function +# --------------------------------------------------------------------------- + +def train_dynamic( + arch: str = "rnn", + train_data: str = "data/dynamic/dynamic_codebert_search.json", + val_data: Optional[str] = None, + val_split: float = 0.05, + save_dir: str = "checkpoints/dynamic", + epochs: int = 10, + batch_size: int = 64, + lr: float = 1e-3, + warmup_ratio: float = 0.1, + max_grad_norm: float = 1.0, + log_every: int = 20, + device_str: Optional[str] = None, + # Model hyperparameters (forwarded to the model constructor) + vocab_size: int = 50265, + num_categories: int = 28, + **model_kwargs, +) -> str: + """Train a dynamic scoring model and return the best checkpoint path. + + Args: + arch: Architecture: ``'linear'``, ``'cnn'``, or ``'rnn'``. + train_data: Path to the JSON file produced by the dataset generator. + val_data: Optional separate validation JSON. When absent, a + fraction ``val_split`` of ``train_data`` is held out. + val_split: Fraction of training data to use as validation when + ``val_data`` is not provided. + save_dir: Directory to save per-epoch and best checkpoints. + epochs: Number of training epochs. + batch_size: Mini-batch size. + lr: Peak learning rate for AdamW. + warmup_ratio: Fraction of total steps used for linear warmup. + max_grad_norm: Gradient clipping threshold. + log_every: Print loss every N batches. + device_str: Force a specific device string (e.g. ``'cpu'``). + vocab_size: Vocabulary size forwarded to the model constructor. + num_categories: Number of AST categories forwarded to the model + constructor. + **model_kwargs: Any additional keyword arguments forwarded to the + selected model constructor (e.g. ``hidden_dim``, + ``kernel_size``). + + Returns: + Path to the best (lowest validation MSE) checkpoint file. + + Raises: + ValueError: For unknown ``arch`` values. + """ + if arch not in _ARCH_MAP: + raise ValueError(f"Unknown arch: {arch!r}. Choose from {list(_ARCH_MAP)!r}.") + + # ── Device ─────────────────────────────────────────────────────────────── + if device_str: + device = torch.device(device_str) + elif torch.cuda.is_available(): + device = torch.device("cuda") + elif torch.backends.mps.is_available(): + device = torch.device("mps") + else: + device = torch.device("cpu") + + # ── Datasets ───────────────────────────────────────────────────────────── + full_ds = DynamicScoringDataset(train_data) + + if val_data is not None: + train_ds = full_ds + val_ds = DynamicScoringDataset(val_data) + else: + val_size = max(1, int(len(full_ds) * val_split)) + train_size = len(full_ds) - val_size + train_ds, val_ds = random_split( + full_ds, [train_size, val_size], + generator=torch.Generator().manual_seed(42), + ) + + train_loader = DataLoader( + train_ds, batch_size=batch_size, shuffle=True, + collate_fn=dynamic_collate_fn, num_workers=0, + ) + val_loader = DataLoader( + val_ds, batch_size=batch_size, shuffle=False, + collate_fn=dynamic_collate_fn, num_workers=0, + ) + + # ── Model ──────────────────────────────────────────────────────────────── + model_cls = _ARCH_MAP[arch] + if arch == "linear": + model = model_cls( + vocab_size=vocab_size, + num_categories=num_categories, + **model_kwargs, + ) + else: + model = model_cls(vocab_size=vocab_size, **model_kwargs) + model.to(device) + + # ── Optimiser / scheduler ───────────────────────────────────────────── + total_steps = len(train_loader) * epochs + warmup_steps = int(total_steps * warmup_ratio) + optimizer = AdamW(model.parameters(), lr=lr, weight_decay=1e-2) + scheduler = get_linear_schedule_with_warmup(optimizer, warmup_steps, total_steps) + + # ── Save dir ───────────────────────────────────────────────────────────── + save_path = Path(save_dir) + save_path.mkdir(parents=True, exist_ok=True) + + print("\n" + "=" * 70) + print(f"TRAINING: Dynamic Scoring Model ({arch.upper()})") + print("=" * 70) + print(f" Architecture : {arch}") + print(f" Train size : {len(train_ds):,}") + print(f" Val size : {len(val_ds):,}") + print(f" Device : {device}") + print(f" Epochs : {epochs}") + print(f" Batch size : {batch_size}") + print(f" LR : {lr}") + print(f" Total steps : {total_steps:,}") + print(f" Save dir : {save_dir}") + print("=" * 70 + "\n") + + best_val_loss = float("inf") + best_ckpt_path = str(save_path / "best.pt") + + for epoch in range(1, epochs + 1): + train_loss = _run_epoch( + model, train_loader, optimizer, scheduler, + arch=arch, device=device, train=True, + max_grad_norm=max_grad_norm, + ) + val_loss = _run_epoch( + model, val_loader, None, None, + arch=arch, device=device, train=False, + ) + + print( + f"Epoch {epoch}/{epochs} | " + f"train_mse={train_loss:.6f} | val_mse={val_loss:.6f}" + ) + + # Save per-epoch checkpoint + epoch_ckpt = str(save_path / f"epoch_{epoch}.pt") + torch.save( + { + "epoch": epoch, + "arch": arch, + "vocab_size": vocab_size, + "num_categories": num_categories, + "model_kwargs": model_kwargs, + "state_dict": model.state_dict(), + "val_loss": val_loss, + }, + epoch_ckpt, + ) + + if val_loss < best_val_loss: + best_val_loss = val_loss + torch.save( + { + "epoch": epoch, + "arch": arch, + "vocab_size": vocab_size, + "num_categories": num_categories, + "model_kwargs": model_kwargs, + "state_dict": model.state_dict(), + "val_loss": val_loss, + }, + best_ckpt_path, + ) + print(f" ↳ new best checkpoint saved (val_mse={val_loss:.6f})") + + print("\nTraining complete.") + print(f"Best checkpoint → {best_ckpt_path} (val_mse={best_val_loss:.6f})\n") + return best_ckpt_path + + +# --------------------------------------------------------------------------- +# Model loading helper +# --------------------------------------------------------------------------- + +def load_dynamic_model(checkpoint_path: str, device_str: Optional[str] = None) -> nn.Module: + """Load a dynamic scoring model from a checkpoint produced by + :func:`train_dynamic`. + + Args: + checkpoint_path: Path to a ``.pt`` file saved during training. + device_str: Optional target device string. + + Returns: + The model in evaluation mode on the requested device. + """ + if device_str: + device = torch.device(device_str) + elif torch.cuda.is_available(): + device = torch.device("cuda") + elif torch.backends.mps.is_available(): + device = torch.device("mps") + else: + device = torch.device("cpu") + + ckpt = torch.load(checkpoint_path, map_location=device) + arch = ckpt["arch"] + vocab_size = ckpt.get("vocab_size", 50265) + num_categories = ckpt.get("num_categories", 28) + model_kwargs = ckpt.get("model_kwargs", {}) + + model_cls = _ARCH_MAP[arch] + if arch == "linear": + model = model_cls( + vocab_size=vocab_size, + num_categories=num_categories, + **model_kwargs, + ) + else: + model = model_cls(vocab_size=vocab_size, **model_kwargs) + + model.load_state_dict(ckpt["state_dict"]) + model.to(device).eval() + return model + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def _build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description="Train a dynamic importance-scoring model") + p.add_argument("--arch", choices=list(_ARCH_MAP), default="rnn") + p.add_argument("--train-data", required=True, + help="Path to JSON produced by dynamic_dataset_generator.py") + p.add_argument("--val-data", default=None, + help="Optional separate validation JSON file.") + p.add_argument("--val-split", type=float, default=0.05) + p.add_argument("--save-dir", default="checkpoints/dynamic") + p.add_argument("--epochs", type=int, default=10) + p.add_argument("--batch-size", type=int, default=64) + p.add_argument("--lr", type=float, default=1e-3) + p.add_argument("--warmup-ratio", type=float, default=0.1) + p.add_argument("--max-grad-norm", type=float, default=1.0) + p.add_argument("--log-every", type=int, default=20) + p.add_argument("--device", default=None) + p.add_argument("--vocab-size", type=int, default=50265) + p.add_argument("--num-categories", type=int, default=28) + # Architecture-specific hyperparameters + p.add_argument("--hidden-dim", type=int, default=128) + p.add_argument("--embed-dim", type=int, default=64) + p.add_argument("--num-layers", type=int, default=2) + p.add_argument("--kernel-size", type=int, default=5) + p.add_argument("--dropout", type=float, default=0.1) + return p + + +def main() -> None: + """CLI entry point.""" + args = _build_arg_parser().parse_args() + + arch_kwargs: dict = {} + if args.arch == "linear": + arch_kwargs = { + "hidden_dim": args.hidden_dim, + "token_embed_dim": args.embed_dim, + "dropout": args.dropout, + } + elif args.arch == "cnn": + arch_kwargs = { + "embed_dim": args.embed_dim, + "hidden_dim": args.hidden_dim, + "num_layers": args.num_layers, + "kernel_size": args.kernel_size, + "dropout": args.dropout, + } + else: # rnn + arch_kwargs = { + "embed_dim": args.embed_dim, + "hidden_dim": args.hidden_dim, + "num_layers": args.num_layers, + "dropout": args.dropout, + } + + train_dynamic( + arch = args.arch, + train_data = args.train_data, + val_data = args.val_data, + val_split = args.val_split, + save_dir = args.save_dir, + epochs = args.epochs, + batch_size = args.batch_size, + lr = args.lr, + warmup_ratio = args.warmup_ratio, + max_grad_norm = args.max_grad_norm, + log_every = args.log_every, + device_str = args.device, + vocab_size = args.vocab_size, + num_categories = args.num_categories, + **arch_kwargs, + ) + + +if __name__ == "__main__": + main() From 571a7deff0ddaf051643f48d32ba9b7af423eaeb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 12 Mar 2026 17:17:24 +0000 Subject: [PATCH 3/4] Address review feedback: use wrapper classes, reuse evaluator logic, CNN lengths param Co-authored-by: hanyiliu <35943468+hanyiliu@users.noreply.github.com> --- src/data/dynamic/dynamic_dataset_generator.py | 200 +++++------------- .../code_summarization_evaluator.py | 17 +- src/evaluation/dynamic/dynamic_evaluator.py | 157 +++++--------- src/models/dynamic/cnn_model.py | 30 ++- src/models/integrated/integrated_codebert.py | 2 +- src/models/integrated/integrated_codet5.py | 2 +- src/training/dynamic/train_dynamic.py | 2 +- 7 files changed, 148 insertions(+), 262 deletions(-) diff --git a/src/data/dynamic/dynamic_dataset_generator.py b/src/data/dynamic/dynamic_dataset_generator.py index 1837322..6867f8d 100644 --- a/src/data/dynamic/dynamic_dataset_generator.py +++ b/src/data/dynamic/dynamic_dataset_generator.py @@ -54,22 +54,22 @@ import json import os from pathlib import Path -from typing import cast +from typing import Any, cast import torch from torch.utils.data import DataLoader from tqdm import tqdm from transformers.models.roberta import RobertaTokenizer -from transformers.models.roberta.modeling_roberta import RobertaModel -from transformers.models.t5 import T5ForConditionalGeneration from src.data.data_loader import load_codesearchnet from src.attention.attention_extraction import ( + extract_batch_cls, + extract_batch_enc_dec, map_subword_tokens_to_categories, _collate, - _cls_scores_from_outputs, - SPECIAL_TOKEN_STRINGS, ) +from src.models.finetuned_codebert_model import FinetunedCodeBERT +from src.models.finetuned_codet5_model import FineTunedCodeT5 # --------------------------------------------------------------------------- # Category vocabulary (statement-level, matching attention_extraction.py) @@ -114,176 +114,87 @@ # --------------------------------------------------------------------------- -# Model loading (from fine-tuned checkpoint) +# Model loading (via fine-tuned wrapper classes) # --------------------------------------------------------------------------- def load_finetuned_codebert( checkpoint_dir: str, device: torch.device, -) -> tuple[RobertaTokenizer, RobertaModel]: - """Load a fine-tuned CodeBERT checkpoint with attention outputs enabled. +) -> tuple[Any, Any]: + """Load a fine-tuned CodeBERT checkpoint using ``FinetunedCodeBERT``. + + Returns the backbone model and tokenizer in eval mode so that + ``extract_batch_cls`` from ``attention_extraction`` can be called + directly on them. Args: checkpoint_dir: Path produced by ``FinetunedCodeBERT.save_model()``. device: Target device. Returns: - (tokenizer, model) — model is in eval mode with attention outputs. + (tokenizer, backbone_model) — model is in eval mode. """ print(f"Loading fine-tuned CodeBERT from '{checkpoint_dir}'...") - tokenizer = cast( - RobertaTokenizer, - RobertaTokenizer.from_pretrained(checkpoint_dir), - ) - model = cast( - RobertaModel, - RobertaModel.from_pretrained(checkpoint_dir, output_attentions=True), - ) - cast(torch.nn.Module, model).to(device).eval() - return tokenizer, model + wrapper = FinetunedCodeBERT(device=str(device)) + wrapper.load_checkpoint(checkpoint_dir) + backbone, tokenizer = wrapper.get_backbone_and_tokenizer() + cast(torch.nn.Module, backbone).to(device).eval() + return tokenizer, backbone def load_finetuned_codet5( checkpoint_dir: str, device: torch.device, -) -> tuple[RobertaTokenizer, T5ForConditionalGeneration]: - """Load a fine-tuned CodeT5 checkpoint with attention outputs enabled. +) -> tuple[Any, Any]: + """Load a fine-tuned CodeT5 checkpoint using ``FineTunedCodeT5``. - Falls back gracefully if the tokenizer metadata uses an unexpected - ``extra_special_tokens`` field (older HF versions). + Returns the backbone model and tokenizer in eval mode so that + ``extract_batch_enc_dec`` from ``attention_extraction`` can be called + directly on them. Args: checkpoint_dir: Path produced by ``FineTunedCodeT5.save_model()``. device: Target device. Returns: - (tokenizer, model) — model is in eval mode with attention outputs. + (tokenizer, backbone_model) — model is in eval mode. """ print(f"Loading fine-tuned CodeT5 from '{checkpoint_dir}'...") - try: - tokenizer = cast( - RobertaTokenizer, - RobertaTokenizer.from_pretrained(checkpoint_dir), - ) - except TypeError as exc: - if "extra_special_tokens" not in str(exc): - raise - tokenizer = cast( - RobertaTokenizer, - RobertaTokenizer.from_pretrained(checkpoint_dir, extra_special_tokens=[]), - ) - model = cast( - T5ForConditionalGeneration, - T5ForConditionalGeneration.from_pretrained(checkpoint_dir, output_attentions=True), - ) - cast(torch.nn.Module, model).to(device).eval() - return tokenizer, model + wrapper = FineTunedCodeT5(device=str(device)) + wrapper.load_checkpoint(checkpoint_dir) + backbone, tokenizer = wrapper.get_backbone_and_tokenizer() + cast(torch.nn.Module, backbone).to(device).eval() + return tokenizer, backbone # --------------------------------------------------------------------------- -# Batched score extraction +# Enrichment: add token_ids / positions / category_ids to attention output # --------------------------------------------------------------------------- -@torch.no_grad() -def _extract_cls_scores( +def _enrich_batch_output( batch: list[dict], - tokenizer: RobertaTokenizer, - model: RobertaModel, - device: torch.device, + tokenizer: Any, + batch_result: list[tuple[list[str], list[float]]], ) -> list[tuple[list[int], list[int], list[int], list[float]]]: - """Extract CLS attention scores for a batch (CodeBERT code-search). - - Returns a list of (token_ids, positions, category_ids, scores) per sample. - Padding tokens are stripped so all lists reflect only the real sequence. - """ - codes = [s.get("func_code_string", "") for s in batch] - enc = tokenizer( - codes, - return_tensors="pt", - truncation=True, - max_length=MAX_LENGTH, - padding=True, - ).to(device) - - outputs = model(**enc, output_attentions=True) - raw_scores = _cls_scores_from_outputs(outputs, len(batch)) - - results = [] - for i, code in enumerate(codes): - real_len = int(enc["attention_mask"][i].sum().item()) - ids_tensor = enc["input_ids"][i, :real_len] - token_ids = ids_tensor.cpu().tolist() - positions = list(range(real_len)) - scores = raw_scores[i][:real_len] - - tokens = tokenizer.convert_ids_to_tokens(ids_tensor) - categories = map_subword_tokens_to_categories(code, tokens) - category_ids = [ - CATEGORY_TO_ID.get(cat, CATEGORY_TO_ID["Other"]) - for cat in categories - ] - results.append((token_ids, positions, category_ids, scores)) - return results + """Convert (tokens, scores) pairs from attention_extraction into the richer + format required for dynamic model training by adding token IDs, positions, + and AST category IDs. + Args: + batch: Raw sample dicts from the DataLoader. + tokenizer: Tokenizer used during extraction (for token → id conversion). + batch_result: Output of ``extract_batch_cls`` or ``extract_batch_enc_dec``; + each element is ``(token_strings, scores)``. -@torch.no_grad() -def _extract_cross_attn_scores( - batch: list[dict], - tokenizer: RobertaTokenizer, - model: T5ForConditionalGeneration, - device: torch.device, -) -> list[tuple[list[int], list[int], list[int], list[float]]]: - """Extract encoder-decoder cross-attention scores (CodeT5 summarization). - - Uses teacher-forcing with the reference description as the decoder input, - matching the approach in ``attention_extraction.extract_batch_enc_dec``. - Returns a list of (token_ids, positions, category_ids, scores) per sample. + Returns: + List of ``(token_ids, positions, category_ids, scores)`` tuples, + one per sample. """ - codes = [s.get("func_code_string", "") for s in batch] - descriptions = [s.get("func_documentation_string", "") for s in batch] - - enc = tokenizer( - codes, - return_tensors="pt", - truncation=True, - max_length=MAX_LENGTH, - padding=True, - ).to(device) - - has_desc = any(descriptions) - if has_desc: - dec_enc = tokenizer( - descriptions, - return_tensors="pt", - truncation=True, - max_length=128, - padding=True, - ).to(device) - decoder_input_ids = dec_enc["input_ids"] - else: - bos_id = tokenizer.pad_token_id or tokenizer.eos_token_id or 0 - decoder_input_ids = torch.full((len(batch), 1), bos_id, device=device) - - outputs = model( - input_ids=enc["input_ids"], - attention_mask=enc["attention_mask"], - decoder_input_ids=decoder_input_ids, - output_attentions=True, - ) - - # cross_attentions[-1]: [B, heads, tgt, src] → avg heads → max over tgt - cross_avg = outputs.cross_attentions[-1].mean(dim=1) # [B, tgt, src] - max_scores = cross_avg.max(dim=1).values # [B, src] - results = [] - for i, code in enumerate(codes): - real_len = int(enc["attention_mask"][i].sum().item()) - ids_tensor = enc["input_ids"][i, :real_len] - token_ids = ids_tensor.cpu().tolist() - positions = list(range(real_len)) - scores = max_scores[i, :real_len].cpu().tolist() - - tokens = tokenizer.convert_ids_to_tokens(ids_tensor) + for sample, (tokens, scores) in zip(batch, batch_result): + code = sample.get("func_code_string", "") + token_ids = tokenizer.convert_tokens_to_ids(tokens) + positions = list(range(len(tokens))) categories = map_subword_tokens_to_categories(code, tokens) category_ids = [ CATEGORY_TO_ID.get(cat, CATEGORY_TO_ID["Other"]) @@ -374,19 +285,10 @@ def generate_dataset( for batch in tqdm(loader, desc="Generating"): try: if extractor == "cls": - batch_out = _extract_cls_scores( - batch, - tokenizer, - cast(RobertaModel, model), - device, - ) + raw_out = extract_batch_cls(batch, tokenizer, model, device, is_t5=False) else: - batch_out = _extract_cross_attn_scores( - batch, - tokenizer, - cast(T5ForConditionalGeneration, model), - device, - ) + raw_out = extract_batch_enc_dec(batch, tokenizer, model, device) + batch_out = _enrich_batch_output(batch, tokenizer, raw_out) except Exception as exc: errors += len(batch) if errors <= 10: diff --git a/src/evaluation/code_summarization_evaluator.py b/src/evaluation/code_summarization_evaluator.py index 91d859e..d062c30 100644 --- a/src/evaluation/code_summarization_evaluator.py +++ b/src/evaluation/code_summarization_evaluator.py @@ -82,6 +82,9 @@ class CodeSummarizationEvaluator: max_src_len : Max encoder input length passed to forward_generation(). max_new_tokens : Max decoder output tokens passed to forward_generation(). num_beams : Beam search width (4 matches LEANCODE paper). + model : Optional pre-constructed ``FineTunedCodeT5`` (or subclass) + instance to use directly. When provided, ``checkpoint_dir`` + and ``device`` are ignored for model construction. Example:: @@ -97,17 +100,21 @@ def __init__( max_src_len: int = 512, max_new_tokens: int = 64, num_beams: int = 4, + model: Optional[FineTunedCodeT5] = None, ): self.max_src_len = max_src_len self.max_new_tokens = max_new_tokens self.num_beams = num_beams - self._model = FineTunedCodeT5(device=device) - - if checkpoint_dir is not None: - self._model.load_checkpoint(checkpoint_dir) + if model is not None: + self._model = model else: - print("[CodeSummarizationEvaluator] No checkpoint — using pretrained weights.") + self._model = FineTunedCodeT5(device=device) + + if checkpoint_dir is not None: + self._model.load_checkpoint(checkpoint_dir) + else: + print("[CodeSummarizationEvaluator] No checkpoint — using pretrained weights.") # ------------------------------------------------------------------ diff --git a/src/evaluation/dynamic/dynamic_evaluator.py b/src/evaluation/dynamic/dynamic_evaluator.py index 1252c33..5ec2e4b 100644 --- a/src/evaluation/dynamic/dynamic_evaluator.py +++ b/src/evaluation/dynamic/dynamic_evaluator.py @@ -64,10 +64,16 @@ from typing import Any, Optional import torch -from nltk.translate.bleu_score import SmoothingFunction, corpus_bleu from tqdm import tqdm from src.data.data_loader import load_codesearchnet +from src.evaluation.code_search_evaluator import ( + _iter_samples, + _score_pairs, + POOL_SIZE, + EVAL_SEED, +) +from src.evaluation.code_summarization_evaluator import CodeSummarizationEvaluator from src.models.integrated.integrated_codebert import IntegratedCodeBERT from src.models.integrated.integrated_codet5 import IntegratedCodeT5 @@ -75,32 +81,7 @@ # Constants # ───────────────────────────────────────────────────────────────────────────── -MAX_LENGTH = 512 -POOL_SIZE = 1000 -EVAL_SEED = 42 -RATIOS = [0.10, 0.20, 0.30, 0.40, 0.50] - - -# ───────────────────────────────────────────────────────────────────────────── -# Dataset normalisation helpers -# ───────────────────────────────────────────────────────────────────────────── - -def _iter_samples(dataset) -> list[dict]: - """Normalise a HuggingFace dataset or JSON list to flat dicts. - - Each dict has keys ``'code'`` and ``'query'`` / ``'reference'``. - """ - samples: list[dict] = [] - for item in dataset: - if isinstance(item, dict): - code = item.get("simplified_code") or item.get("func_code_string", "") - query = item.get("description") or item.get("func_documentation_string", "") - else: - code = item["func_code_string"] - query = item["func_documentation_string"] - if code.strip() and query.strip(): - samples.append({"code": code, "query": query, "reference": query}) - return samples +RATIOS = [0.10, 0.20, 0.30, 0.40, 0.50] # ───────────────────────────────────────────────────────────────────────────── @@ -111,6 +92,10 @@ class DynamicEvaluator: """Evaluates an integrated dynamic model on code search (MRR) or code summarisation (BLEU-4). + Reuses ``_iter_samples`` and ``_score_pairs`` from + ``code_search_evaluator`` and delegates BLEU-4 computation to + ``CodeSummarizationEvaluator`` with the integrated model injected. + Args: task: ``'search'`` or ``'summarization'``. model_type: ``'codebert'`` or ``'codet5'``. @@ -171,58 +156,33 @@ def __init__( ) if model_type == "codebert": - self._model: Any = IntegratedCodeBERT( + self._integrated_model: Any = IntegratedCodeBERT( dynamic_checkpoint = dynamic_checkpoint, codebert_checkpoint = downstream_checkpoint, simplification_ratio = simplification_ratio, device = device, ) else: - self._model = IntegratedCodeT5( + self._integrated_model = IntegratedCodeT5( dynamic_checkpoint = dynamic_checkpoint, codet5_checkpoint = downstream_checkpoint, simplification_ratio = simplification_ratio, device = device, ) - _, tokenizer = self._model._load_model() + _, tokenizer = self._integrated_model._load_model() self._tokenizer: Any = tokenizer - # ── Code-search MRR ────────────────────────────────────────────────────── - - @torch.no_grad() - def _score_pairs( - self, - query: str, - codes: list[str], - ) -> torch.Tensor: - """Score (query, code) pairs with class-1 logit after simplification.""" - all_scores: list[torch.Tensor] = [] - - # Simplify each code snippet before scoring - simplified_codes = [self._model.simplify_code(c) or c for c in codes] - - for start in range(0, len(simplified_codes), self.batch_size): - chunk = simplified_codes[start : start + self.batch_size] - query_chunk = [query] * len(chunk) - enc = self._tokenizer( - query_chunk, - chunk, - return_tensors = "pt", - truncation = True, - max_length = MAX_LENGTH, - padding = True, - ).to(self.device) - - logits = self._model.forward_classification( - input_ids = enc["input_ids"], - attention_mask = enc["attention_mask"], + # For summarization, build a CodeSummarizationEvaluator with the integrated + # model injected directly. IntegratedCodeT5.forward_generation already + # applies dynamic simplification before generating, so compute_bleu + # works without any further changes. + if task == "summarization": + self._summ_evaluator = CodeSummarizationEvaluator( + model = self._integrated_model, ) - if isinstance(logits, tuple): - logits = logits[0] - all_scores.append(logits[:, 1].detach().cpu()) - return torch.cat(all_scores, dim=0) + # ── Code-search MRR ────────────────────────────────────────────────────── def compute_mrr( self, @@ -233,6 +193,11 @@ def compute_mrr( ) -> tuple[float, float]: """Compute MRR on dataset with dynamic simplification. + Uses ``_iter_samples`` and ``_score_pairs`` from + ``code_search_evaluator``. Before scoring, each candidate code is + dynamically simplified by the integrated model's ``simplify_code`` + method. + Args: dataset: HuggingFace Dataset or list[dict]. pool_size: Candidate pool per query (1 correct + negatives). @@ -254,15 +219,30 @@ def compute_mrr( reciprocal_ranks = [] for i, sample in enumerate(tqdm(samples, desc="MRR (dynamic)")): - query = sample["query"] - neg_indices = rng.sample( + query = sample["query"] + neg_indices = rng.sample( [j for j in range(n) if j != i], k=min(pool_size - 1, n - 1), ) - pool_indices = [i] + neg_indices - pool_codes = [samples[j]["code"] for j in pool_indices] + pool_indices = [i] + neg_indices + pool_codes_raw = [samples[j]["code"] for j in pool_indices] + + # Dynamically simplify each candidate before scoring + pool_codes = [ + self._integrated_model.simplify_code(c) or c + for c in pool_codes_raw + ] + + # Reuse _score_pairs from code_search_evaluator for the core scoring + scores = _score_pairs( + query = query, + codes = pool_codes, + tokenizer = self._tokenizer, + model = self._integrated_model, + device = self.device, + batch_size = self.batch_size, + ) - scores = self._score_pairs(query, pool_codes) ranked = scores.argsort(descending=True).tolist() rank_of_correct = ranked.index(0) + 1 reciprocal_ranks.append(1.0 / rank_of_correct) @@ -281,6 +261,11 @@ def compute_bleu( ) -> tuple[float, float]: """Compute corpus BLEU-4 on dataset with dynamic simplification. + Delegates entirely to ``CodeSummarizationEvaluator.compute_bleu`` + with the integrated model injected. ``IntegratedCodeT5.forward_generation`` + already applies dynamic simplification before generating, so no + additional pre-processing is needed here. + Args: dataset: HuggingFace Dataset or list[dict]. max_samples: Evaluate on at most this many examples. @@ -289,39 +274,9 @@ def compute_bleu( Returns: ``(bleu4, elapsed_seconds)`` """ - samples = _iter_samples(dataset) - if max_samples is not None: - samples = samples[:max_samples] - n = len(samples) - if n == 0: - print("[DynamicEvaluator] Empty dataset — returning 0.") - return 0.0, 0.0 - - print(f"\n[DynamicEvaluator] BLEU-4 evaluation on {n:,} samples...") - - hypotheses: list[list[str]] = [] - references: list[list[list[str]]] = [] - t_start = time.time() - - for i, sample in enumerate(samples): - try: - pred = self._model.forward_generation(sample["code"]) - except Exception as exc: - print(f" [WARNING] sample {i} failed: {exc}") - pred = "" - - hypotheses.append(pred.split()) - references.append([sample["reference"].split()]) - - if (i + 1) % log_every == 0 or (i + 1) == n: - print(f" Progress: {i + 1:,}/{n:,}") - - elapsed = time.time() - t_start - sf = SmoothingFunction().method4 - bleu_100 = corpus_bleu(references, hypotheses, smoothing_function=sf) * 100.0 - - print(f" Corpus BLEU-4: {bleu_100:.2f} | elapsed: {elapsed:.1f}s") - return bleu_100, elapsed + return self._summ_evaluator.compute_bleu( + dataset, max_samples=max_samples, log_every=log_every + ) # ── Unified evaluate entry point ───────────────────────────────────────── diff --git a/src/models/dynamic/cnn_model.py b/src/models/dynamic/cnn_model.py index 3741488..cdd9b58 100644 --- a/src/models/dynamic/cnn_model.py +++ b/src/models/dynamic/cnn_model.py @@ -27,11 +27,16 @@ model = CNNScoringModel(kernel_size=5) # token_ids: LongTensor [batch, seq_len] - scores = model(token_ids) # FloatTensor [batch, seq_len] + scores = model(token_ids) # FloatTensor [batch, seq_len] + + # With sequence lengths to zero out padding positions: + scores = model(token_ids, lengths=lengths) # FloatTensor [batch, seq_len] """ from __future__ import annotations +from typing import Optional + import torch import torch.nn as nn @@ -51,7 +56,8 @@ class CNNScoringModel(nn.Module): Usage:: model = CNNScoringModel(kernel_size=5) - scores = model(token_ids) # [batch, seq_len] + scores = model(token_ids) # [batch, seq_len] + scores = model(token_ids, lengths) # zeroes out padding positions """ def __init__( @@ -88,11 +94,19 @@ def __init__( self.conv_stack = nn.Sequential(*conv_layers) - def forward(self, token_ids: torch.Tensor) -> torch.Tensor: + def forward( + self, + token_ids: torch.Tensor, + lengths: Optional[torch.Tensor] = None, + ) -> torch.Tensor: """Compute importance scores for a batch of token sequences. Args: token_ids: LongTensor of token IDs with shape ``[batch, seq_len]``. + lengths: Optional LongTensor ``[batch]`` with the actual + (non-padded) length of each sequence. When supplied, + scores at padding positions are zeroed out, consistent + with the RNN model's behaviour. Returns: FloatTensor of per-token scores in ``(0, 1)``, @@ -101,4 +115,12 @@ def forward(self, token_ids: torch.Tensor) -> torch.Tensor: # [B, T, embed_dim] → transpose to [B, embed_dim, T] for Conv1d x = self.embedding(token_ids).transpose(1, 2) out = self.conv_stack(x) # [B, 1, T] - return out.squeeze(1) # [B, T] + scores = out.squeeze(1) # [B, T] + + if lengths is not None: + # Build a mask: True for valid positions, False for padding + B, T = scores.shape + mask = torch.arange(T, device=scores.device).unsqueeze(0) < lengths.unsqueeze(1) + scores = scores * mask.float() + + return scores diff --git a/src/models/integrated/integrated_codebert.py b/src/models/integrated/integrated_codebert.py index ed422f5..018ee58 100644 --- a/src/models/integrated/integrated_codebert.py +++ b/src/models/integrated/integrated_codebert.py @@ -113,7 +113,7 @@ def _simplify_code( elif arch == "rnn": scores = dynamic_model(token_ids_t, lengths_t).squeeze(0).tolist() else: # cnn - scores = dynamic_model(token_ids_t).squeeze(0).tolist() + scores = dynamic_model(token_ids_t, lengths_t).squeeze(0).tolist() # Identify removable positions (non-special tokens only) removable_indices = [ diff --git a/src/models/integrated/integrated_codet5.py b/src/models/integrated/integrated_codet5.py index cf64c33..d7a885e 100644 --- a/src/models/integrated/integrated_codet5.py +++ b/src/models/integrated/integrated_codet5.py @@ -110,7 +110,7 @@ def _simplify_code( elif arch == "rnn": scores = dynamic_model(token_ids_t, lengths_t).squeeze(0).tolist() else: # cnn - scores = dynamic_model(token_ids_t).squeeze(0).tolist() + scores = dynamic_model(token_ids_t, lengths_t).squeeze(0).tolist() # Identify removable positions (non-special tokens only) removable_indices = [ diff --git a/src/training/dynamic/train_dynamic.py b/src/training/dynamic/train_dynamic.py index 3c87542..d926b2d 100644 --- a/src/training/dynamic/train_dynamic.py +++ b/src/training/dynamic/train_dynamic.py @@ -127,7 +127,7 @@ def _run_epoch( if arch == "linear": preds = model(token_ids, positions, category_ids) elif arch == "cnn": - preds = model(token_ids) + preds = model(token_ids, lengths) else: # rnn preds = model(token_ids, lengths) From 6d37e06e9bcf71ed9ab86f114d1058df38702501 Mon Sep 17 00:00:00 2001 From: Hanyi Liu Date: Fri, 24 Apr 2026 16:31:11 -0700 Subject: [PATCH 4/4] Update dynamic_approaches.md --- context/dynamic_approaches.md | 1 + 1 file changed, 1 insertion(+) diff --git a/context/dynamic_approaches.md b/context/dynamic_approaches.md index 53100be..131c7f4 100644 --- a/context/dynamic_approaches.md +++ b/context/dynamic_approaches.md @@ -10,6 +10,7 @@ Static methods' primary weakness is that they cannot address the importance of c Therefore, to address the context uniqueness of all code snippets, we must approach code simplification in a dynamic method that can capture code context and dynamically assign importance score to each token before inference. A perfect dynamic method would satisfy the following: 1. Input: Tokenized code snippet. Output: Importance score for each individual token. 2. The compute required for this importance calculation must be much smaller than passing the code snippet through a multi-layer transformer model. +3. The input code snippet would be passed through this method's model first to calculate a simplified version, then this would be passed into the actual LLM. The dynamic method would then be some form of a neural network model. To train this potential model, we build off of LeanCode's assumption that the attention score of the last transformer layer is most indicative of a token's importance. With this in mind, we then essentially have free training data. For all code snippets, we can pass it through our finetuned transformer models, extract their attention scores, and use the (code snippet, attention scores) pair as training data for our dynamic model.