naos-ku/GraphTokenLM

1from typing import Literal2 3import torch4import torch.nn as nn5from torch_geometric.nn import (6    GATConv,7    GCNConv,8    GINConv,9    GraphSAGE,10    TransformerConv,11    global_add_pool,12    global_max_pool,13    global_mean_pool,14)15from torch_geometric.utils import to_dense_batch16from transformers import (17    AutoConfig,18    AutoModelForCausalLM,19    PretrainedConfig,20    PreTrainedModel,21)22from transformers.generation.utils import GenerationMixin23from transformers.modeling_outputs import CausalLMOutputWithPast24 25VALID_GRAPH_POOLING = ["mean", "sum", "max"]26 27 28class GraphTokenLMConfig(PretrainedConfig):29    model_type = "graph_token_lm"30 31    def __init__(32        self,33        base_model="Qwen/Qwen3-4B-Base",34        gnn_type: Literal["GCN", "GAT", "GIN", "GraphSAGE", "GraphTransformer"] = "GCN",35        node_feat_dim=8,36        lpe_dim: int | None = None,37        use_degree_emb: bool = False,38        pos_emb_dim=8,39        gnn_hidden_dim=256,40        gnn_out_dim=512,41        num_gnn_layers=2,42        graph_pooling: list[Literal["mean", "sum", "max"]] = ["mean"],43        num_proj_layers=1,44        num_graph_tokens=4,45        num_max_nodes=20,  # maximum number of nodes per batch46        freeze_llm=True,47        enable_lora: bool = False,48        tie_word_embeddings=True,49        **kwargs,50    ):51        """Initialize a GraphToken language model configuration.52 53        Parameters54        ----------55        base_model : str, default="Qwen/Qwen3-4B-Base"56            Hugging Face model name or path for the underlying LLM.57        gnn_type : {"GCN", "GAT", "GIN", "GraphSAGE", "GraphTransformer"}, default="GCN"58            Type of GNN layer to use for encoding graph nodes.59        node_feat_dim : int, default=860            Dimensionality of the raw node features.61        lpe_dim : int or None, default=None62            Dimensionality for Laplacian positional encodings; defaults to ``node_feat_dim``.63        use_degree_emb : bool, default=False64            Whether to add degree-based embeddings to node features.65        pos_emb_dim : int, default=866            Dimensionality of learned positional embeddings for nodes.67        gnn_hidden_dim : int, default=25668            Hidden dimensionality for the GNN stack.69        gnn_out_dim : int, default=51270            Output dimensionality of the GNN encoder.71        num_gnn_layers : int, default=272            Number of GNN layers to apply.73        num_proj_layers : int, default=174            Number of projection layers mapping graph reps to tokens.75        num_graph_tokens : int, default=476            Number of graph tokens to prepend to the LLM.77        num_max_nodes : int, default=2078            Maximum number of nodes per graph in a batch.79        graph_pooling : {"mean", "sum", "max"} list, default=["mean"]80            Pooling strategy for graph-level aggregation. When multiple values are81            provided, pooled vectors are concatenated.82        freeze_llm : bool, default=True83            Whether to freeze the underlying LLM parameters.84        enable_lora : bool, default=False85            Whether LoRA adapters are expected to be active, in which case the86            base LLM should remain in training mode unless explicitly set87            elsewhere.88        tie_word_embeddings : bool, default=True89            Whether to tie input/output embeddings in the LLM config.90        **kwargs91            Additional arguments forwarded to ``PretrainedConfig``.92        """93        self.base_model = base_model94        self.llm_name = base_model  # backward compatibility95        self.gnn_type = gnn_type96 97        self.node_feat_dim = node_feat_dim98        self.lpe_dim = lpe_dim if lpe_dim is not None else node_feat_dim99        self.use_degree_emb = bool(use_degree_emb)100        self.pos_emb_dim = pos_emb_dim101        self.node_pos_emb_dim = pos_emb_dim  # backward compatibility102        self.gnn_hidden_dim = gnn_hidden_dim103        self.gnn_hidden = gnn_hidden_dim  # backward compatibility104        self.gnn_out_dim = gnn_out_dim105        self.gnn_out = gnn_out_dim  # backward compatibility106        self.num_gnn_layers = num_gnn_layers107        self.num_proj_layers = num_proj_layers108        self.num_graph_tokens = num_graph_tokens109        self.num_max_nodes = num_max_nodes110        self.graph_pooling = graph_pooling111        self.freeze_llm = freeze_llm112        self.enable_lora = enable_lora113 114        # Keep generation-related fields for compatibility (updated later).115        self.vocab_size = kwargs.get("vocab_size", None)116        self.pad_token_id = kwargs.get("pad_token_id", None)117        self.bos_token_id = kwargs.get("bos_token_id", None)118        self.eos_token_id = kwargs.get("eos_token_id", None)119 120        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)121 122    # Maintain compatibility with transformers.GenerationConfig.123    def get_text_config(self, decoder: bool | None = None, **kwargs):124        return self125 126 127class GNNEncoder(nn.Module):128 129    def __init__(130        self,131        in_dim: int,132        hid_dim: int,133        out_dim: int,134        max_nodes: int,135        num_layers: int = 2,136        node_pos_emb_dim: int = 8,137        dropout: float = 0.1,138        gnn_type: Literal["GCN", "GAT", "GIN", "GraphSAGE", "GraphTransformer"] = "GCN",139    ):140        """Encode graph node features with a configurable GNN stack.141 142        Parameters143        ----------144        in_dim : int145            Dimensionality of the input node features.146        hid_dim : int147            Hidden dimensionality used for intermediate layers.148        out_dim : int149            Dimensionality of the output node representations.150        max_nodes : int151            Maximum number of nodes per graph in a batch.152        num_layers : int, default=2153            Number of graph convolution layers.154        node_pos_emb_dim : int, default=8155            Dimensionality of the optional learned positional embeddings.156        dropout : float, default=0.1157            Dropout probability applied between hidden layers.158        gnn_type : {"GCN", "GAT", "GIN", "GraphSAGE", "GraphTransformer"}, default="GCN"159            Type of graph convolution layer to build.160        """161        super().__init__()162        self.max_nodes = max_nodes163        self.pos_emb = nn.Embedding(max_nodes, node_pos_emb_dim) if node_pos_emb_dim > 0 else None164 165        in_channels = in_dim + (node_pos_emb_dim if node_pos_emb_dim > 0 else 0)166        if in_channels <= 0:167            raise ValueError("GNNEncoder requires a positive input feature dimension.")168 169        hidden_dims = [hid_dim] * max(num_layers - 1, 0)170        dims = [in_channels, *hidden_dims, out_dim]171        self.convs = nn.ModuleList()172        for i in range(len(dims) - 1):173            match gnn_type:174                case "GCN":175                    self.convs.append(GCNConv(dims[i], dims[i + 1]))176                case "GAT":177                    self.convs.append(GATConv(dims[i], dims[i + 1]))178                case "GIN":179                    self.convs.append(GINConv(nn.Linear(dims[i], dims[i + 1])))180                case "GraphSAGE":181                    self.convs.append(GraphSAGE(dims[i], dims[i + 1], 1))182                case "GraphTransformer":183                    # Multi-head attention with concat disabled to keep the output dim aligned.184                    self.convs.append(TransformerConv(dims[i], dims[i + 1], heads=4, concat=False, dropout=dropout))185                case _:186                    raise ValueError(f"Unsupported gnn_type: {gnn_type}")187        self.act = nn.ReLU()188        self.dropout = nn.Dropout(dropout)189 190    def forward(self, x, edge_index, batch):191        if self.pos_emb is not None:192            # Assign positional indices per graph starting from zero within the batch.193            _, mask = to_dense_batch(x, batch, max_num_nodes=self.max_nodes)194            pos_idx = torch.arange(self.max_nodes, device=x.device).unsqueeze(0).expand(mask.size(0), -1)195            pos_idx = pos_idx[mask]196            x = torch.cat([x, self.pos_emb(pos_idx)], dim=-1)197        for i, conv in enumerate(self.convs):198            x = conv(x, edge_index)199            if i < len(self.convs) - 1:200                x = self.act(x)201                x = self.dropout(x)202        return x  # [num_nodes, out_dim]203 204 205def _normalize_graph_pooling(graph_pooling: list[str]) -> list[str]:206    """Normalize and validate graph pooling values into a canonical list."""207    poolings = list(graph_pooling)208    if not (1 <= len(poolings) <= 3):209        raise ValueError(f"Unsupported graph_pooling length: {len(poolings)}")210    if len(set(poolings)) != len(poolings):211        raise ValueError(f"Duplicate graph_pooling values: {poolings}")212    if not set(poolings).issubset(set(VALID_GRAPH_POOLING)):213        raise ValueError(f"Unsupported graph_pooling values: {poolings}")214    return poolings215 216 217class DomainProjector(nn.Module):218 219    def __init__(220        self,221        gnn_out_dim,222        llm_hidden_size,223        num_graph_tokens=4,224        num_layers=1,225        graph_pooling: list[str] = ["mean"],226    ):227        """Project graph-level representations into graph tokens.228 229        The projector first pools node embeddings into a graph representation and230        then maps it into ``k`` graph tokens that match the language model's hidden231        dimension.232 233        Parameters234        ----------235        gnn_out_dim : int236            Dimensionality of the encoder output to project from.237        llm_hidden_size : int238            Target dimensionality matching the language model embeddings.239        num_graph_tokens : int, default=4240            Number of graph tokens to produce.241        graph_pooling : {"mean", "sum", "max"} list, default=["mean"]242            Pooling strategy used to aggregate node embeddings. When multiple values are243            provided, pooled vectors are concatenated.244        num_layers : int, default=1245            Number of linear/GELU projection layers.246        """247        super().__init__()248        if num_layers < 1:249            raise ValueError("DomainProjector requires at least one projection layer.")250        poolings = _normalize_graph_pooling(graph_pooling)251        self.num_graph_tokens = num_graph_tokens252        self.graph_pooling = poolings253        layers = []254        in_dim = gnn_out_dim * len(poolings)255        final_dim = llm_hidden_size * num_graph_tokens256        if num_layers == 1:257            layer_dims = [final_dim]258        else:259            ratio = final_dim / in_dim260            layer_dims = []261            prev_dim = in_dim262            for layer_idx in range(num_layers):263                t = (layer_idx + 1) / num_layers264                dim = max(1, int(round(in_dim * (ratio ** t))))265                if dim % 2 != 0:266                    dim += 1267                dim = max(dim, prev_dim)268                layer_dims.append(dim)269                prev_dim = dim270            layer_dims[-1] = final_dim271        for layer_idx, out_dim in enumerate(layer_dims):272            layers.append(nn.Linear(in_dim, out_dim))273            if layer_idx < num_layers - 1:274                layers.append(nn.GELU())275            in_dim = out_dim276        self.project = nn.Sequential(*layers)277 278        # Optional learned positional embeddings for graph tokens.279        self.graph_pos = nn.Embedding(num_graph_tokens, llm_hidden_size)280 281    def forward(self, node_repr, batch_index):282        """Aggregate node embeddings into graph tokens.283 284        Parameters285        ----------286        node_repr : torch.Tensor287            Node representations of shape ``(num_nodes_total, gnn_out_dim)``.288        batch_index : torch.Tensor289            Batch indices identifying the graph for each node. Shape290            ``(num_nodes_total,)``.291 292        Returns293        -------294        torch.Tensor295            Graph token tensor of shape ``(batch_size, num_graph_tokens, hidden)``.296        """297        # Pool node representations into a graph-level vector.298        pooled_list = []299        if "mean" in self.graph_pooling:300            pooled_list.append(global_mean_pool(node_repr, batch_index))  # [B, gnn_out_dim]301        if "sum" in self.graph_pooling:302            pooled_list.append(global_add_pool(node_repr, batch_index))  # [B, gnn_out_dim]303        if "max" in self.graph_pooling:304            pooled_list.append(global_max_pool(node_repr, batch_index))  # [B, gnn_out_dim]305        pooled = (306            torch.cat(pooled_list, dim=-1) if len(pooled_list) > 1 else pooled_list[0]307        )  # [B, gnn_out_dim * num_poolings]308        B = pooled.size(0)309 310        # Expand into k tokens via the projection stack.311        tokens = self.project(pooled)  # [B, k * hidden]312        Hk = tokens.view(B, self.num_graph_tokens, -1)  # [B, k, hidden]313 314        # Add learned positional embeddings.315        pos = self.graph_pos.weight.unsqueeze(0).expand(B, -1, -1)  # [B, k, hidden]316        Hk = Hk + pos317        return Hk  # [B, k, hidden]318 319 320class GraphTokenLM(PreTrainedModel, GenerationMixin):321    """Language model that prepends graph tokens to textual inputs.322 323    A graph neural network encodes node features, pools them, and projects the324    result into learned graph tokens that are concatenated with language model325    embeddings before decoding.326 327    Parameters328    ----------329    config : GraphTokenLMConfig330        Model configuration describing the graph encoder and base LLM.331    load_llm_weights : bool, default=True332        Whether to load pretrained weights for the base language model.333    """334 335    _tied_weights_keys = ["llm.lm_head.weight"]336    _keys_to_ignore_on_load_missing = [r"^llm\.lm_head\.weight$"]337 338    config_class = GraphTokenLMConfig339    base_model_prefix = "llm"340 341    def __init__(self, config: GraphTokenLMConfig, load_llm_weights: bool = True):342 343        super().__init__(config)344 345        # LLM346        if load_llm_weights:347            self.llm = AutoModelForCausalLM.from_pretrained(348                config.base_model, trust_remote_code=True, tie_word_embeddings=True349            )350        else:351            llm_cfg = AutoConfig.from_pretrained(config.base_model, dtype=torch.float32)352            self.llm = AutoModelForCausalLM.from_config(llm_cfg)353 354        self.num_graph_tokens = config.num_graph_tokens355 356        # GNN + Domain Projector357        self.gnn = GNNEncoder(358            gnn_type=config.gnn_type,359            node_pos_emb_dim=config.pos_emb_dim,360            in_dim=config.node_feat_dim,361            hid_dim=config.gnn_hidden_dim,362            out_dim=config.gnn_out_dim,363            num_layers=config.num_gnn_layers,364            max_nodes=config.num_max_nodes,365        )366        self.tokenizer_head = DomainProjector(367            gnn_out_dim=config.gnn_out_dim,368            llm_hidden_size=self.llm.config.hidden_size,369            num_graph_tokens=config.num_graph_tokens,370            num_layers=config.num_proj_layers,371            graph_pooling=config.graph_pooling,372        )373 374        if config.freeze_llm:375            for p in self.llm.parameters():376                p.requires_grad = False377            if not config.enable_lora:378                self.llm.eval()379 380        # --- sync basic generation fields so GenerationMixin works cleanly ---381        mirror_keys = [382            "vocab_size",383            "pad_token_id",384            "bos_token_id",385            "eos_token_id",386            "hidden_size",387            "num_hidden_layers",388            "num_attention_heads",389        ]390        for k in mirror_keys:391            if hasattr(self.llm.config, k):392                setattr(self.config, k, getattr(self.llm.config, k))393 394        # make sure tying is done once at init (harmless if already tied)395        if getattr(self.config, "tie_word_embeddings", False):396            self.tie_weights()397 398    @property399    def device(self):400        return next(self.parameters()).device401 402    def _concat_graph_tokens(403        self,404        input_ids=None,405        attention_mask=None,406        labels=None,407        inputs_embeds=None,408        graph=None,409    ):410        """Prepend graph tokens to language model embeddings.411 412        Parameters413        ----------414        input_ids : torch.Tensor, optional415            Token IDs used to derive embeddings if ``inputs_embeds`` is not provided.416        attention_mask : torch.Tensor, optional417            Attention mask aligned with ``input_ids``.418        labels : torch.Tensor, optional419            Label tensor passed through unchanged.420        inputs_embeds : torch.Tensor, optional421            Precomputed language model embeddings.422        graph : Mapping[str, torch.Tensor], optional423            Graph structure containing ``x``, ``edge_index``, and ``batch`` as424            produced by the collator.425 426        Returns427        -------428        tuple of torch.Tensor429            Tuple ``(inputs_embeds, attention_mask, labels)`` with graph tokens430            concatenated at the front of the sequence.431        """432        # Obtain embeddings from the base language model if necessary.433        if inputs_embeds is None:434            inputs_embeds = self.llm.get_input_embeddings()(input_ids)435 436        B, T, H = inputs_embeds.size()437 438        # ---- Graph to tokens ----439        graph_device = next(self.gnn.parameters()).device440        if hasattr(graph, "to"):441            graph = graph.to(graph_device)442 443        x = graph["x"]  # [N_nodes, node_feat_dim]444        edge_index = graph["edge_index"]  # [2, N_edges]445        batch = graph["batch"]  # [N_nodes]446        node_repr = self.gnn(x, edge_index, batch)  # [N_nodes, gnn_out]447        graph_tokens = self.tokenizer_head(node_repr, batch)  # [B, k, H]448        if inputs_embeds is not None:449            graph_tokens = graph_tokens.to(inputs_embeds.device)450 451        # ---- Concatenate by prepending graph tokens ----452        new_inputs = torch.cat([graph_tokens, inputs_embeds], dim=1)  # [B, k+T, H]453 454        # Prepend ones to the attention mask for the graph tokens.455        if attention_mask is None:456            attention_mask = input_ids.ne(self.llm.config.pad_token_id).long()457        new_attention = torch.cat(458            [torch.ones((B, self.num_graph_tokens), dtype=attention_mask.dtype, device=self.device), attention_mask],459            dim=1,460        )461 462        return new_inputs, new_attention, labels463 464    def forward(465        self,466        input_ids=None,467        attention_mask=None,468        inputs_embeds=None,469        graph=None,470        labels=None,471        **generate_kwargs,472    ) -> CausalLMOutputWithPast:473        if (input_ids is None) and (inputs_embeds is None):474            raise ValueError("Either input_ids or inputs_embeds must be provided")475 476        if graph is not None:477            inputs_embeds, attention_mask, labels = self._concat_graph_tokens(478                input_ids=input_ids,479                attention_mask=attention_mask,480                labels=labels,481                inputs_embeds=inputs_embeds,482                graph=graph,483            )484        else:485            # For generation, inputs may already include concatenated graph tokens.486            if inputs_embeds is None:487                inputs_embeds = self.llm.get_input_embeddings()(input_ids)488            if attention_mask is None:489                if input_ids is None:490                    raise ValueError("When attention_mask is not provided, input_ids must also be provided")491                attention_mask = input_ids.ne(self.llm.config.pad_token_id).long()492 493        # Exclude auxiliary keys that may interfere with Qwen3 loss computation.494        blocked = {495            "num_items_in_batch",496            "label_smoothing",  # Added by TRL/transformers in some setups.497            "labels_shifted",  # Same as above.498        }499        passdown = {500            k: v501            for k, v in generate_kwargs.items()502            if k not in blocked and k in {"use_cache", "output_attentions", "output_hidden_states", "past_key_values"}503        }504 505        out = self.llm(506            inputs_embeds=inputs_embeds,507            attention_mask=attention_mask,508            labels=labels,509            **passdown,510        )511        return out512 513    def prepare_inputs_for_generation(514        self, input_ids=None, inputs_embeds=None, attention_mask=None, graph=None, **kwargs515    ):516        if input_ids is None and inputs_embeds is None:517            raise ValueError("Either input_ids or inputs_embeds must be provided")518        if input_ids is not None and inputs_embeds is not None:519            raise ValueError("Both input_ids and inputs_embeds cannot be provided at the same time")520 521        if inputs_embeds is None:522            inputs_embeds = self.llm.get_input_embeddings()(input_ids)523        if graph is not None:524            inputs_embeds, attention_mask, _ = self._concat_graph_tokens(525                input_ids=None,526                attention_mask=attention_mask,527                labels=None,528                inputs_embeds=inputs_embeds,529                graph=graph,530            )531 532        return {"inputs_embeds": inputs_embeds, "attention_mask": attention_mask, "graph": None}533 534    # delegate embeddings to inner LLM so HF can tie weights correctly535    def get_input_embeddings(self):536        return self.llm.get_input_embeddings()537 538    def set_input_embeddings(self, new_embeddings):539        self.llm.set_input_embeddings(new_embeddings)540 541    def get_output_embeddings(self):542        return self.llm.get_output_embeddings()543 544    def set_output_embeddings(self, new_embeddings):545        self.llm.set_output_embeddings(new_embeddings)546 547    def tie_weights(self):548        # honor config.tie_word_embeddings and delegate549        if getattr(self.config, "tie_word_embeddings", False):550            # inner LLM handles actual tying (lm_head <-> embeddings)551            self.llm.tie_weights()552        # keep parent behavior (no-op for most models)553        return super().tie_weights()554