lstm-poem-generator

UNILSTM Poem Generator

A character-level UNILSTM model fine-tuned for poetry generation.

Model Description

This is a custom PyTorch LSTM model trained on a poetry dataset for generating poems.

Architecture

• Model Type: unilstm
• Embedding Dimension: 256
• Hidden Dimension: 512
• Number of Layers: 3
• Dropout: 0.3
• Max Length: 512

Training

• Batch Size: 32
• Learning Rate: 0.001
• Epochs: 50
• Optimizer: AdamW
• Scheduler: CosineAnnealingLR

Usage

import torch
from model import BiLSTMPoetryModel  # or UnidirectionalLSTMPoetryModel
from shared.preprocess import CharVocabulary

# Load vocabulary
vocab = CharVocabulary.load("vocab.json")

# Load model
model = BiLSTMPoetryModel(
    vocab_size=vocab.vocab_size,
    embedding_dim=256,
    hidden_dim=512,
    num_layers=3,
    dropout=0.0,  # No dropout for inference
    pad_idx=vocab.pad_idx
)

checkpoint = torch.load("best_model.pt", map_location="cpu")
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Generate a poem
prompt = "Title: The Moon\n\n"
start_tokens = vocab.encode(prompt, add_special_tokens=False)
generated = model.generate(start_tokens, vocab, max_length=300, device="cpu")
print(generated)

Files

• best_model.pt: PyTorch model checkpoint
• vocab.json: Character vocabulary mapping
• config.json: Training configuration

License

MIT License