import json
import torch

from safetensors.torch import load_file
from transformers import AutoTokenizer

from model import GPTModel,generate_and_print_sample

# load config
with open("config.json") as f:
    cfg = json.load(f)

# create model
model = GPTModel(cfg)
model.to("cuda")

# load weights
state_dict = load_file("model.safetensors")

model.load_state_dict(state_dict)

model.eval()

# tokenizer
tokenizer = AutoTokenizer.from_pretrained(".")

print(generate_and_print_sample(model, tokenizer, "cuda", "The world is big"))