from transformers import pipeline

# Initialize the text generation pipeline with the distilgpt2 model
generator = pipeline('text-generation', model='distilgpt2', device=0)  # Use GPU (device 0)


# Generate text with the specified parameters
output = generator(
    'in this course, we will teach how to',
    max_length=20,  # Set maximum length of the generated text
    num_return_sequences=1  # Generate 2 sequences
)

# Print the full output
print(output)

# Print each generated sequence separately
for i, result in enumerate(output):
    print(f"Sequence {i + 1}: {result['generated_text']}")