from transformers import GPT2Tokenizer, GPT2LMHeadModel

# Load the fine-tuned model
model_path = "./fine_tuned_model"
model = GPT2LMHeadModel.from_pretrained(model_path)
tokenizer = GPT2Tokenizer.from_pretrained(model_path)

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 with attention mask
        attention_mask = input_ids.ne(tokenizer.pad_token_id)
        output = model.generate(input_ids, max_length=100, num_return_sequences=1, temperature=0.7, do_sample=True, attention_mask=attention_mask)

        # 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))