aframson/RDPDLM

1import torch2import torch.nn as nn3import torch.nn.functional as F4from transformers.modeling_utils import PreTrainedModel5 6# Define your custom language model class7class OBILanguageModel(PreTrainedModel):8    def __init__(self, config):9        super(OBILanguageModel,self).__init__(config)10        self.token_embedding_table = nn.Embedding(config.vocab_size, config.hidden_size)  # Use length of SentencePiece vocab11        self.position_embedding_table = nn.Embedding(config.block_size, config.hidden_size)12        self.transformer = nn.Transformer(13            d_model=config.hidden_size,14            nhead=config.num_attention_heads,15            num_encoder_layers=config.num_hidden_layers,16            num_decoder_layers=config.num_hidden_layers,17            dim_feedforward=4 * config.hidden_size,18            dropout=config.hidden_dropout_prob,19            activation='gelu'20        )21        self.ln1 = nn.LayerNorm(config.hidden_size)22        self.ln2 = nn.LayerNorm(config.hidden_size)23        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)  # Use length of SentencePiece vocab24    25    def forward(self, idx, targets=None):26        tok_emb = self.token_embedding_table(idx)27        pos_emb = self.position_embedding_table(torch.arange(idx.size(1), device='cpu'))28        x = tok_emb + pos_emb29        x = self.transformer(x, x)30        x = self.ln1(x)31        x = self.ln2(x)32        logits = self.lm_head(x)33        34        if targets is None:35            loss = None36        else:37            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), targets.view(-1))38        39        return logits, loss40 41    def generate(self, idx, max_new_tokens):42        for _ in range(max_new_tokens):43            idx_cond = idx[:, -self.config.block_size:]44            logits, loss = self(idx_cond)45            logits = logits[:, -1, :]46            probs = F.softmax(logits, dim=-1)47            idx_next = torch.multinomial(probs, num_samples=1)48            idx = torch.cat((idx, idx_next), dim=1)49        return idx50