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

kriteekathapa/lstm-poem-generator-v1

UNILSTM Poem Generator

Model Description

Architecture

Training

Usage

Files

License

kriteekathapa/
lstm-poem-generator-v1