import argparse
from transformers import GPTNeoForCausalLM, GPT2Tokenizer # Add GPTNeoForCausalLM import here
import torch

def generate_text(input_text):
# Load pre-trained model and tokenizer
model_name = "EleutherAI/gpt-neo-1.3B" # Specify model name
model = GPTNeoForCausalLM.from_pretrained(model_name) # Load GPT-Neo model
tokenizer = GPT2Tokenizer.from_pretrained(model_name) # Load tokenizer for GPT-Neo

# Ensure that the pad_token_id is set to the eos_token_id for GPT-2 models
model.config.pad_token_id = model.config.eos_token_id

# Encode the input prompt
inputs = tokenizer.encode(input_text, return_tensors="pt")

# Create the attention mask (all ones, since there's no padding involved)
attention_mask = torch.ones(inputs.shape, device=inputs.device)

# Generate text from the model with adjustments for more sensible output
output = model.generate(
    inputs, 
    attention_mask=attention_mask, 
    max_length=100,  # Set max_length to a reasonable value
    num_return_sequences=1,  # Only generate one sequence
    do_sample=True,  # Enable sampling for some variety
    temperature=0.6,  # Lower temperature for more deterministic output
    top_p=0.85,  # Nucleus sampling with a smaller top_p for more focused output
    top_k=30,  # Limit to top 30 tokens for sampling
    pad_token_id=model.config.pad_token_id  # Ensure pad_token_id is set correctly
)

# Decode the output
decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)

# Print the result
print(decoded_output)

if name == "main":
# Set up argument parsing
parser = argparse.ArgumentParser(description="Generate text with GPT-2 or GPT-Neo")
parser.add_argument("input", type=str, help="Input text prompt for GPT-2")

# Parse the arguments
args = parser.parse_args()

# Call the generate_text function with the user-provided input
generate_text(args.input)