File size: 2,545 Bytes
4798617 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | import torch
from transformers import AutoTokenizer
from torch import nn
device1 = torch.device("cuda:0")
device2 = torch.device("cuda:1")
class SplitModel(nn.Module):
def __init__(self, embedding_layer, dropout_layer, gptj_blocks, layer_norm, lm_head):
super(SplitModel, self).__init__()
self.embedding_layer = embedding_layer
self.dropout_layer = dropout_layer
self.gptj_blocks = gptj_blocks
self.layer_norm = layer_norm
self.lm_head = lm_head
def forward(self, input_ids, attention_mask):
tensor_ids = self.dropout_layer(self.embedding_layer(input_ids))
# GPTJBlock is missing the embedding positions that are necessary for self-attention.
# To fix this issue, you need to ensure that the position_ids are passed to each GPTJBlock during the forward pass.
position_ids = torch.arange(tensor_ids.shape[1], dtype=torch.long, device=tensor_ids.device)
for block in self.gptj_blocks:
tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
tensor_ids = tensor_ids.to(device2)
tensor_ids = self.layer_norm(tensor_ids)
logits = self.lm_head(tensor_ids)
logits = logits.to(device1)
return logits
model_dir = "pt_fp32"
model_path = f"{model_dir}/torch_model.pt"
tokenizer = AutoTokenizer.from_pretrained(model_dir)
full_model = torch.load(model_path)
embedding_layer = full_model.transformer.wte.to(device1)
dropout_layer = full_model.transformer.drop.to(device1)
gptj_blocks = full_model.transformer.h.to(device1)
layer_norm = full_model.transformer.ln_f.to(device2)
lm_head = full_model.lm_head.to(device2)
split_model = SplitModel(embedding_layer, dropout_layer, gptj_blocks, layer_norm, lm_head)
input_text = "Hi I am Jade and I love"
input("Press enter please")
input_tokens = tokenizer.encode_plus(input_text, return_tensors="pt").to(device1)
input_ids = input_tokens["input_ids"]
temperature = 0.8
max_new_tokens = 50
with torch.no_grad():
for _ in range(max_new_tokens):
attention_mask = torch.ones_like(input_ids).to(device1)
logits = split_model(input_ids, attention_mask)[:, -1] / temperature
probabilities = torch.softmax(logits, dim=-1)
sampled_token_ids = torch.multinomial(probabilities, num_samples=1)
input_ids = torch.cat((input_ids, sampled_token_ids), dim=-1)
generated_ids = input_ids.squeeze().tolist()
output = tokenizer.decode(generated_ids, skip_special_tokens=True)
print(output) |