import argparse
from transformers import GPT2LMHeadModel, GPT2Tokenizer


def test_model(prompt, model_path='/Users/raghul.v/Desktop/research/pii_extraction_test/results', model_name='distilgpt2'):  # Adjust the default model_name
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_path)  # Loading the fine-tuned model

    # Note: Adjustments for running the model on a CUDA device are commented out,
    # Uncomment and use those if CUDA is available and desired.

    # Encode the prompt to test
    generated = tokenizer.encode(prompt, return_tensors="pt")

    # Generate text from the model
    sample_outputs = model.generate(
        generated,
        do_sample=True,
        max_length=50,
        top_k=50,
        top_p=0.95,
        num_return_sequences=3
    )

    for idx, sample_output in enumerate(sample_outputs):
        decoded_output = tokenizer.decode(sample_output, skip_special_tokens=True)
        print(f"Generated Text {idx}: {decoded_output}")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Enter the prompt for the model.")
    parser.add_argument('--prompt', type=str, required=True, help='Prompt for the model')
    args = parser.parse_args()

    # Default model_path is 'results', adjust if the model is saved elsewhere
    test_model(args.prompt, model_path='results', model_name='distilgpt2')