from transformers import GPT2Tokenizer, GPT2LMHeadModel

# Load the original GPT-2 tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

# Load your fine-tuned model
model = GPT2LMHeadModel.from_pretrained("./fine_tuned_model")

# Save the tokenizer along with your fine-tuned model
model.save_pretrained("./fine_tuned_model")
tokenizer.save_pretrained("./fine_tuned_model")

while True:
    try:
        # Prompt user for input
        prompt_text = input("You: ")

        # Tokenize the prompt text
        input_ids = tokenizer.encode(prompt_text, return_tensors="pt")

        # Generate response
        output = model.generate(input_ids, max_length=100, num_return_sequences=1, temperature=0.7, do_sample=True)

        # Decode and print the generated response
        generated_response = tokenizer.decode(output[0], skip_special_tokens=True)
        print("Bot:", generated_response)

    except Exception as e:
        print("An error occurred:", str(e))