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)