poem-generator-lstm / train_model.py
Sukumar2005's picture
Update train_model.py
6d728c6 verified
# train_model.py
import numpy as np
from keras.models import Sequential
from keras.layers import LSTM, Dense
from keras.utils import to_categorical
from keras.callbacks import ModelCheckpoint
# Load and preprocess data
with open("poem_data.txt", "r", encoding='utf-8') as file:
text = file.read().lower()
chars = sorted(list(set(text)))
char_to_idx = {c: i for i, c in enumerate(chars)}
idx_to_char = {i: c for i, c in enumerate(chars)}
seq_length = 40
step = 1
X = []
y = []
for i in range(0, len(text) - seq_length, step):
seq = text[i:i + seq_length]
label = text[i + seq_length]
X.append([char_to_idx[c] for c in seq])
y.append(char_to_idx[label])
X = np.array(X)
y = to_categorical(y, num_classes=len(chars))
# Build the model
model = Sequential()
model.add(LSTM(128, input_shape=(seq_length, 1)))
model.add(Dense(len(chars), activation='softmax'))
model.compile(loss='categorical_crossentropy', optimizer='adam')
# Reshape and normalize
X = np.reshape(X, (X.shape[0], seq_length, 1)) / float(len(chars))
# Train the model
model.fit(X, y, batch_size=128, epochs=20)
# Save model
model.save("model.keras")