import torch
import os
from gpt_dev import GPTLanguageModel, encode, decode, generate_text  # Import necessary parts from gpt_dev.py

# Set up device (GPU or CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load your saved model (pre-trained model's state dict)
block_size= 256
model_type = "gpt"
n_embd = 384
n_head = 6
n_layer = 6
transformers_version = 4.44.2
vocab_size = 95

# Instantiate the model
model = GPTLanguageModel(vocab_size, embedding_size, num_heads, num_layers)
model.to(device)

# Load the pre-trained weights
if os.path.exists("gpt_language_model.pth"):
    checkpoint = torch.load("gpt_language_model.pth", map_location=device)
    model.load_state_dict(checkpoint)

# Generate text based on a prompt
start_prompt = "Once upon a time"
generated_text = generate_text(model, start_prompt, max_length=100)
print(generated_text)  # Display the generated text