model.py

"""
Model Definition for SVG Generation

Defines the components of the SVG generation model, including the GPT-2 based
decoder, word token embeddings, and the final language model head.
"""

from torch import nn
from transformers import GPT2Config, GPT2Model

def build_svg_decoder(vocab_size: int, d_model: int, n_head: int, bos_token_id: int, eos_token_id: int):
    """
    Builds the SVG decoder components.

    Args:
        vocab_size (int): The size of the vocabulary.
        d_model (int): The dimensionality of the model.
        n_head (int): The number of attention heads.
        bos_token_id (int): The ID of the beginning-of-sequence token.
        eos_token_id (int): The ID of the end-of-sequence token.

    Returns:
        tuple: A tuple containing the SVG decoder (GPT2Model),
               the token embedding layer (nn.Embedding),
               and the language model head (nn.Linear).
    """
    config = GPT2Config(
        vocab_size=vocab_size,
        n_embd=d_model,
        n_head=n_head,
        bos_token_id=bos_token_id,
        eos_token_id=eos_token_id,
        add_cross_attention=True
    )

    svg_decoder = GPT2Model(config)
    wte = nn.Embedding(config.vocab_size, d_model)
    lm_head = nn.Linear(d_model, vocab_size, bias=False)
    
    return svg_decoder, wte, lm_head