import os
from groq import Groq

client = Groq(api_key='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')

def generate_poems(num_poems, length_per_poem, output_dir):
    # Create output directory if it doesn't exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    for i in range(num_poems):
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a talented poet."
                },
                {
                    "role": "user",
                    "content": f"Write a complex poem with a length of {length_per_poem} tokens."
                }
            ],
            model="llama3-8b-8192",
            temperature=0.7,
            max_tokens=length_per_poem,
            top_p=1,
            stop=None,
            stream=False,
        )
        
        poem = chat_completion.choices[0].message.content.strip()
        
        # Save the poem to a text file
        with open(os.path.join(output_dir, f'poem_{i + 1}.txt'), 'w') as f:
            f.write(poem)

# Parameters
num_poems = 20  # Number of poems to generate
length_per_poem = 150  # Length of each poem in tokens (adjust accordingly)
output_dir = 'poems'  # Directory to save poems

# Generate and save poems
generate_poems(num_poems, length_per_poem, output_dir)