|
from transformers import AutoConfig, PreTrainedModel, AutoModelForCausalLM
|
|
from typing import List, Optional
|
|
from torch import nn
|
|
from model import LayerNorm, BlockJ
|
|
from transformers.modeling_outputs import CausalLMOutputWithPast
|
|
import torch
|
|
import math
|
|
from torch.nn import functional as F
|
|
from transformers import AutoConfig, AutoModel
|
|
from .pretrained_config import *
|
|
|
|
|
|
|
|
class GPTJXForCausalLM(PreTrainedModel):
|
|
config_class = GPTJXConfig
|
|
base_model_prefix = "transformer"
|
|
is_parallelizable = True
|
|
supports_gradient_checkpointing = True
|
|
_no_split_modules = ["BlockJ"]
|
|
|
|
_supports_flash_attn_2 = True
|
|
_tied_weights_keys = ["lm_head.weight"]
|
|
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
assert config.vocab_size is not None
|
|
assert config.block_size is not None
|
|
self.config = config
|
|
|
|
self.transformer = nn.ModuleDict(dict(
|
|
wte = nn.Embedding(config.vocab_size, config.n_embd),
|
|
wpe = nn.Embedding(config.block_size, config.n_embd),
|
|
drop = nn.Dropout(config.dropout),
|
|
h = nn.ModuleList([BlockJ(config) for _ in range(config.n_layer)]),
|
|
ln_f = LayerNorm(config.n_embd, bias=config.bias),
|
|
))
|
|
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
|
|
|
|
|
|
|
|
|
|
self.transformer.wte.weight = self.lm_head.weight
|
|
|
|
|
|
self.apply(self._init_weights)
|
|
|
|
for pn, p in self.named_parameters():
|
|
if pn.endswith('c_proj.weight'):
|
|
torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
|
|
|
|
|
|
print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))
|
|
|
|
def get_num_params(self, non_embedding=True):
|
|
"""
|
|
Return the number of parameters in the model.
|
|
For non-embedding count (default), the position embeddings get subtracted.
|
|
The token embeddings would too, except due to the parameter sharing these
|
|
params are actually used as weights in the final layer, so we include them.
|
|
"""
|
|
n_params = sum(p.numel() for p in self.parameters())
|
|
if non_embedding:
|
|
n_params -= self.transformer.wpe.weight.numel()
|
|
return n_params
|
|
|
|
|
|
def get_input_embeddings(self):
|
|
return self.wte
|
|
|
|
def set_input_embeddings(self, new_embeddings):
|
|
self.wte = new_embeddings
|
|
|
|
def forward(self, idx, targets=None, attn_mask= None, output_hidden_states: Optional[bool] = None, **kwargs):
|
|
device = idx.device
|
|
b, t = idx.size()
|
|
|
|
|
|
|
|
|
|
assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
|
|
pos = torch.arange(0, t, dtype=torch.long, device=device)
|
|
|
|
|
|
tok_emb = self.transformer.wte(idx)
|
|
pos_emb = self.transformer.wpe(pos)
|
|
x = self.transformer.drop(tok_emb + pos_emb)
|
|
for block in self.transformer.h:
|
|
x = block(x, attn_mask=attn_mask)
|
|
x = self.transformer.ln_f(x)
|
|
|
|
if targets is not None:
|
|
|
|
logits = self.lm_head(x)
|
|
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)
|
|
else:
|
|
|
|
logits = self.lm_head(x[:, [-1], :])
|
|
loss = None
|
|
|
|
|
|
return CausalLMOutputWithPast(
|
|
loss=loss,
|
|
logits=logits,
|
|
hidden_states=x if output_hidden_states else None,
|
|
attentions= None,
|
|
)
|
|
|
|
def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
|
|
|
|
model_inputs = {"idx": input_ids}
|
|
|
|
|
|
if attention_mask is not None:
|
|
model_inputs["attn_mask"] = attention_mask
|
|
|
|
return model_inputs
|
|
|
|
|
|
@torch.no_grad()
|
|
def stream(self, idx, max_new_tokens, temperature=1.0, top_k=None,gen_mode="greedy"):
|
|
"""
|
|
Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
|
|
the sequence max_new_tokens times, feeding the predictions back into the model each time.
|
|
Most likely you'll want to make sure to be in model.eval() mode of operation for this.
|
|
"""
|
|
for _ in range(max_new_tokens):
|
|
|
|
idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
|
|
|
|
logits, _ = self(idx_cond, eval=True)
|
|
|
|
logits = logits[:, -1, :] / temperature
|
|
|
|
if top_k is not None:
|
|
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
|
|
logits[logits < v[:, [-1]]] = -float('Inf')
|
|
|
|
probs = F.softmax(logits, dim=-1)
|
|
|
|
if gen_mode == 'greedy':
|
|
idx_next = torch.argmax(probs, dim=-1).unsqueeze(0)
|
|
|
|
else:
|
|
idx_next = torch.multinomial(probs, num_samples=1)
|
|
|
|
idx = torch.cat((idx, idx_next), dim=1)
|
|
|
|
yield idx_next
|
|
|
|
|
|
def crop_block_size(self, block_size):
|
|
|
|
|
|
|
|
assert block_size <= self.config.block_size
|
|
self.config.block_size = block_size
|
|
self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:block_size])
|
|
for block in self.transformer.h:
|
|
if hasattr(block.attn, 'bias'):
|
|
block.attn.bias = block.attn.bias[:,:,:block_size,:block_size]
|
|
|
|
|
|
AutoConfig.register("nanogpt-j", GPTJXConfig)
|
|
AutoModel.register(GPTJXConfig,GPTJXForCausalLM)
|
|
AutoModelForCausalLM.register(GPTJXConfig, GPTJXForCausalLM)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
from transformers import AutoTokenizer
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("BeardedMonster/SabiYarn")
|
|
input_ids = tokenizer("ba wo ni?", return_tensors="pt")["input_ids"]
|
|
targets = input_ids
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
model = GPTJXForCausalLM.from_pretrained("/pretrainedmodel")
|
|
|
|
outputs = model(input_ids, targets)
|
|
print(outputs)
|
|
output = model.generate(input_ids, max_new_tokens=100)
|
|
print(tokenizer.decode(output[0]))
|
|
|