"""GraphCodeBERT clone-detection model. Reimplements the architecture from Microsoft's ``GraphCodeBERT/clonedetection`` on top of ``transformers`` v5: * the two snippets are encoded **separately** by one shared GraphCodeBERT encoder, each with its own graph-guided masked attention; * a data-flow node's input embedding is the average of the embeddings of the code tokens it was identified from; * the two ```` representations are concatenated and fed to a ``Linear(2H -> H) -> tanh -> Linear(H -> 2)`` head. The only real adaptation is the attention mask: ``transformers`` v5 builds masks through ``masking_utils`` and only forwards a mask untouched when it is already 4-D, so the boolean ``[B, L, L]`` graph mask is expanded to an additive ``[B, 1, L, L]`` mask here. """ from __future__ import annotations import torch import torch.nn as nn from transformers import RobertaConfig, RobertaModel, RobertaPreTrainedModel from transformers.modeling_outputs import SequenceClassifierOutput __all__ = ["GraphCodeBERTForCloneDetection", "CloneClassificationHead"] def _autocast_dtype(device: torch.device, fallback: torch.dtype) -> torch.dtype: """Dtype the attention scores will actually have, honouring autocast. SDPA requires an additive ``attn_mask`` whose dtype matches the query, so a hard-coded float32 mask would break under ``fp16=True``. """ try: if torch.is_autocast_enabled(device.type): return torch.get_autocast_dtype(device.type) except TypeError: # older signature without a device argument if device.type == "cuda" and torch.is_autocast_enabled(): return torch.get_autocast_gpu_dtype() return fallback def _to_additive_mask(bool_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: """``[B, L, L]`` boolean -> ``[B, 1, L, L]`` additive mask (0 / -inf).""" additive = torch.zeros(bool_mask.shape, dtype=dtype, device=bool_mask.device) additive.masked_fill_(~bool_mask, torch.finfo(dtype).min) return additive.unsqueeze(1) class CloneClassificationHead(nn.Module): """Pairwise head over the two ```` vectors (GraphCodeBERT's own head).""" def __init__(self, config: RobertaConfig) -> None: super().__init__() self.dense = nn.Linear(config.hidden_size * 2, config.hidden_size) self.dropout = nn.Dropout(config.hidden_dropout_prob) self.out_proj = nn.Linear(config.hidden_size, 2) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: """``hidden_states``: ``[B*2, L, H]`` -> logits ``[B, 2]``.""" x = hidden_states[:, 0, :] # of each snippet x = x.reshape(-1, x.size(-1) * 2) # pair the two snippets back up x = self.dropout(x) x = torch.tanh(self.dense(x)) x = self.dropout(x) return self.out_proj(x) class GraphCodeBERTForCloneDetection(RobertaPreTrainedModel): """Binary clone classifier: ``0 = not clone``, ``1 = clone``.""" config_class = RobertaConfig base_model_prefix = "roberta" supports_gradient_checkpointing = True def __init__(self, config: RobertaConfig, class_weights: list[float] | None = None) -> None: super().__init__(config) config.num_labels = 2 self.roberta = RobertaModel(config, add_pooling_layer=False) self.classifier = CloneClassificationHead(config) self.register_buffer( "class_weights", torch.tensor(class_weights, dtype=torch.float32) if class_weights else None, persistent=False, ) self.post_init() # ------------------------------------------------------------------ # def _embed_with_dataflow( self, input_ids: torch.Tensor, position_idx: torch.Tensor, attn_mask: torch.Tensor ) -> torch.Tensor: """Word embeddings where each data-flow node averages its code tokens. ``position_idx`` encodes the role of every slot: ``0`` = data-flow node, ``1`` (= ````) = padding, ``>= 2`` = real code token. """ nodes_mask = position_idx.eq(0) token_mask = position_idx.ge(2) embeddings = self.roberta.embeddings.word_embeddings(input_ids) # For every node row, the code-token columns it may look at. nodes_to_token = nodes_mask[:, :, None] & token_mask[:, None, :] & attn_mask nodes_to_token = nodes_to_token.to(embeddings.dtype) nodes_to_token = nodes_to_token / (nodes_to_token.sum(-1) + 1e-10)[:, :, None] averaged = torch.einsum("abc,acd->abd", nodes_to_token, embeddings) return embeddings * (~nodes_mask)[:, :, None] + averaged * nodes_mask[:, :, None] def _encode( self, input_ids: torch.Tensor, position_idx: torch.Tensor, attn_mask: torch.Tensor ) -> torch.Tensor: embeddings = self._embed_with_dataflow(input_ids, position_idx, attn_mask) dtype = _autocast_dtype(input_ids.device, embeddings.dtype) outputs = self.roberta( inputs_embeds=embeddings, attention_mask=_to_additive_mask(attn_mask, dtype), position_ids=position_idx, token_type_ids=torch.zeros_like(position_idx), ) return outputs.last_hidden_state # ------------------------------------------------------------------ # def forward( self, input_ids_1: torch.Tensor, position_idx_1: torch.Tensor, attn_mask_1: torch.Tensor, input_ids_2: torch.Tensor, position_idx_2: torch.Tensor, attn_mask_2: torch.Tensor, labels: torch.Tensor | None = None, ) -> SequenceClassifierOutput: """Encode both snippets with the shared encoder and classify the pair. Args: input_ids_*: ``[B, L]`` token ids; data-flow slots hold ````. position_idx_*: ``[B, L]`` role/position ids (see ``_embed_with_dataflow``). attn_mask_*: ``[B, L, L]`` boolean graph-guided attention mask. labels: ``[B]`` with values in ``{0, 1}``. """ batch_size, seq_len = input_ids_1.shape # Stack both snippets into one encoder call: [B, L] x2 -> [B*2, L]. input_ids = torch.cat((input_ids_1[:, None], input_ids_2[:, None]), 1).view(-1, seq_len) position_idx = torch.cat((position_idx_1[:, None], position_idx_2[:, None]), 1).view( -1, seq_len ) attn_mask = torch.cat((attn_mask_1[:, None], attn_mask_2[:, None]), 1).view( -1, seq_len, seq_len ) hidden = self._encode(input_ids, position_idx, attn_mask) logits = self.classifier(hidden) loss = None if labels is not None: weight = None if self.class_weights is not None: weight = self.class_weights.to(device=logits.device, dtype=logits.dtype) loss = nn.functional.cross_entropy(logits, labels.view(-1), weight=weight) return SequenceClassifierOutput(loss=loss, logits=logits) def load_model( model_name_or_path: str, attn_implementation: str = "sdpa", class_weights: list[float] | None = None, gradient_checkpointing: bool = False, ) -> GraphCodeBERTForCloneDetection: """Load GraphCodeBERT weights into the pairwise clone-detection head.""" model = GraphCodeBERTForCloneDetection.from_pretrained( model_name_or_path, attn_implementation=attn_implementation, ) # Set after loading: `from_pretrained` should not have to carry runtime-only # arguments, and the weights are a training artefact, not part of the config. model.class_weights = ( torch.tensor(class_weights, dtype=torch.float32) if class_weights else None ) if model.config.model_type != "roberta": raise ValueError( f"Expected a RoBERTa-architecture checkpoint (GraphCodeBERT), " f"got model_type={model.config.model_type!r}." ) if gradient_checkpointing: model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) return model def count_parameters(model: nn.Module) -> dict[str, int]: trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) total = sum(p.numel() for p in model.parameters()) return {"trainable_parameters": trainable, "total_parameters": total}