from transformers import AutoTokenizer
import torch

device1 = torch.device("cuda:0")
device2 = torch.device("cuda:1")

class SplitModel(torch.nn.Module):
    def __init__(self, base_model):
        super(SplitModel, self).__init__()
        self.embedding_layer = base_model.transformer.wte.to(device1)
        # self.dropout_layer = base_model.transformer.drop.to(device1)
        self.gptj_blocks1 = torch.nn.ModuleList(base_model.transformer.h[:14]).to(device1)
        self.gptj_blocks2 = torch.nn.ModuleList(base_model.transformer.h[14:]).to(device2)
        self.layer_norm = base_model.transformer.ln_f.to(device2)
        self.lm_head = base_model.lm_head.to(device2)
    
    def forward(self, input_ids, attention_mask):
        # tensor_ids = self.dropout_layer(self.embedding_layer(input_ids))
        tensor_ids = self.embedding_layer(input_ids)
        position_ids = torch.arange(tensor_ids.shape[1], dtype=torch.long, device=tensor_ids.device)
        for block in self.gptj_blocks1:
            tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
        tensor_ids = tensor_ids.to(device2)
        position_ids = position_ids.to(device2)
        attention_mask = attention_mask.to(device2)
        for block in self.gptj_blocks2:
            tensor_ids = block(tensor_ids, attention_mask=attention_mask, position_ids=position_ids)[0]
        tensor_ids = self.layer_norm(tensor_ids)
        logits = self.lm_head(tensor_ids)
        return logits.to(device1)

model_dir = "pt_fp32"
model_path = f"{model_dir}/torch_model.pt"

tokenizer = AutoTokenizer.from_pretrained(model_dir)
split_model = SplitModel(torch.load(model_path))

input_text = "Hi I am Jade and I love"
input_tokens = tokenizer.encode_plus(input_text, return_tensors="pt").to(device1)
input_ids = input_tokens["input_ids"]
temperature = 0.5
max_new_tokens = 50
with torch.no_grad():
    # split_model.eval()
    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)
        del logits, probabilities, sampled_token_ids
    generated_ids = input_ids.squeeze().tolist()
output = tokenizer.decode(generated_ids, skip_special_tokens=True)
print(output)