jmssndrlnd/rhymes

1import datetime2import os3import random4from transformers import GPT2TokenizerFast, GPT2LMHeadModel5 6os.environ["TOKENIZERS_PARALLELISM"] = "false"7 8# Initialize GPT-2 model and tokenizer9tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")10model = GPT2LMHeadModel.from_pretrained("gpt2")11 12 13def generate_rhyme(time_str):14    input_text = f"Include '{time_str}' in a short rhyme: "15    input_ids = tokenizer.encode(input_text, return_tensors="pt")16 17    # Use sampling instead of greedy decoding18    outputs = model.generate(19        input_ids,20        max_length=50,21        num_return_sequences=5,22        do_sample=True,23        top_k=50,24        top_p=0.95,25    )26 27    generated_texts = [28        tokenizer.decode(output, skip_special_tokens=True) for output in outputs29    ]30    rhymes = []31 32    for text in generated_texts:33        if input_text in text:34            rhymes.append(text.split(input_text)[1].strip())35        else:36            rhymes.append("")37 38    # Filter rhymes with 20 words or less39    filtered_rhymes = [rhyme for rhyme in rhymes if len(rhyme.split()) <= 20]40 41    if filtered_rhymes:42        return random.choice(filtered_rhymes)43    else:44        return f"At {time_str}, enjoy your day."45 46 47def main():48    with open("rhymes.txt", "w", encoding="utf-8") as output_file:49        for hour in range(24):50            for minute in range(60):51                timestamp = datetime.datetime.strftime(52                    datetime.datetime(100, 1, 1, hour, minute), "%I:%M %p"53                )54                rhyme = generate_rhyme(timestamp)55                output_file.write(f"{timestamp} - {rhyme}\n")56 57 58if __name__ == "__main__":59    main()60