LSTM Next Word Prediction Model

This project trains an LSTM-based language model to predict the next word in a quote. The model is trained on a quote dataset, where each quote is cleaned, tokenized, converted into prefix-target pairs, padded to a fixed sequence length, and then learned as a multiclass next-token classification problem.

The final implementation uses PyTorch for training and inference. Some notebook cells also contain earlier TensorFlow/Keras experiments with SimpleRNN and LSTM models, but the saved checkpoint in this repository belongs to the PyTorch LSTM workflow.

Project Structure

.
|-- README.md
|-- main.py
|-- pyproject.toml
|-- datasets/
|   `-- qoute_dataset.csv
|-- config/
|   |-- X.npy
|   |-- y.npy
|   |-- tokenizer.pkl
|   `-- config.pkl
|-- model/
|   `-- next_word_checkpoint.pth
|-- notebooks/
|   |-- pre-process.ipynb
|   |-- train_model.ipynb
|   `-- test_model.ipynb
`-- screenshots/
    `-- training and output screenshots

Dataset

The dataset is stored at datasets/qoute_dataset.csv.

It contains 3,038 quote records with two columns:

• quote: the quote text used for training
• Author: the author of the quote

Only the quote column is used for the language model. The author labels are not used in the current training pipeline.

Preprocessing

The preprocessing workflow is shown in notebooks/pre-process.ipynb.

The main steps are:

1. Load the quote dataset with Pandas.
2. Select the quote column.
3. Convert all text to lowercase.
4. Remove punctuation using Python's string.punctuation.
5. Fit a Keras Tokenizer on the cleaned quotes.
6. Convert every quote into a sequence of integer word ids.
7. Build next-word training pairs from every quote.
8. Pad all input sequences to a fixed length.
9. Save the processed arrays and tokenizer for reuse.

For a quote sequence like:

life is a journey

the generated training samples follow this pattern:

Input: life              Target: is
Input: life is           Target: a
Input: life is a         Target: journey

After preprocessing, the saved training arrays are:

X shape: (85271, 745)
y shape: (85271,)

The final PyTorch workflow keeps y as integer class labels and trains with nn.CrossEntropyLoss. Earlier notebook experiments briefly used one-hot encoded labels for Keras, but that is not used by the saved PyTorch checkpoint.

Saved Configuration

The model configuration is saved in config/config.pkl.

VOCAB_SIZE = 10000
MAX_LEN = 745
EMBED_DIM = 128
HIDDEN_DIM = 128

The tokenizer is saved at config/tokenizer.pkl.

Tokenizer details from the saved artifact:

Vocabulary cap: 10000
Out-of-vocabulary token: <OOV>
Observed word index size: 8979

Model Architecture

The final model is a single-layer LSTM implemented in PyTorch.

class LSTMModel(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            batch_first=True
        )
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        output, (hidden, cell) = self.lstm(x)
        hidden = hidden[-1]
        out = self.fc(hidden)
        return out

Architecture summary:

Input sequence          : (batch_size, 745)
Embedding               : Embedding(10000, 128)
Embedding output        : (batch_size, 745, 128)
LSTM                    : LSTM(input_size=128, hidden_size=128, batch_first=True)
Final hidden state      : (batch_size, 128)
Fully connected layer   : Linear(128, 10000)
Output logits           : (batch_size, 10000)
Prediction target       : next word id

Approximate trainable parameter count:

Embedding layer         : 1,280,000
LSTM layer              :   132,096
Linear output layer     : 1,290,000
Total                   : 2,702,096

The model outputs raw logits over the vocabulary. During training, these logits are passed directly into nn.CrossEntropyLoss. During inference, the highest-probability token is selected with argmax and appended back to the seed text.

Training

Training is implemented in notebooks/train_model.ipynb.

The training notebook:

• Loads X.npy, y.npy, tokenizer.pkl, and config.pkl
• Builds a PyTorch Dataset and DataLoader
• Trains with batch size 128
• Uses Adam optimizer with learning rate 0.001
• Uses nn.CrossEntropyLoss
• Saves and resumes from next_word_checkpoint.pth
• Uses GPU automatically when CUDA is available

Training settings used in the notebook:

Batch size    : 128
Optimizer     : Adam
Learning rate : 0.001
Loss          : CrossEntropyLoss
Device        : CUDA if available, otherwise CPU
Checkpoint    : next_word_checkpoint.pth

The notebook output shows training on a Tesla T4 GPU and resuming from an existing checkpoint at epoch 48. The final visible checkpointed run completed epoch 148 with:

Average loss : 0.1710
Accuracy     : 95.57%

The checkpoint is stored in this repository at:

model/next_word_checkpoint.pth

The checkpoint contains the model state and optimizer state so training can continue from the last saved epoch.

Inference

Inference is implemented in notebooks/test_model.ipynb.

The inference flow:

1. Load the saved tokenizer.
2. Load the saved model configuration.
3. Recreate the same LSTMModel architecture.
4. Load the checkpoint weights.
5. Accept seed text from the user.
6. Generate words one at a time by repeatedly feeding the growing text back into the model.

Example output from the training notebook:

Seed text:
life is

Generated text:
life is a series of natural and spontaneous changes dont resist them that only creates sorrow let

Example output from the test notebook:

Seed text:
life is

Generated text:
life is pain highness anyone who says differently is selling something in paperback for ever or place but rather they let the

The generated text is produced with greedy decoding, meaning the model always chooses the token with the highest predicted score. This makes inference simple and deterministic, but it can also make the text repetitive or less diverse than sampling-based generation.

Running the Project

Create and activate a Python environment, then install the project dependencies.

pip install -e .

The notebooks use PyTorch, but torch is not currently listed in pyproject.toml. Install a PyTorch build that matches your system before running the training or inference notebooks.

For a CPU-only install:

pip install torch

For CUDA-enabled training, install PyTorch using the command recommended for your CUDA version from the official PyTorch installation page.

After installing dependencies, open the notebooks:

jupyter notebook notebooks/pre-process.ipynb
jupyter notebook notebooks/train_model.ipynb
jupyter notebook notebooks/test_model.ipynb

The notebooks were originally written for a Google Colab or Drive-style working directory where the dataset, config files, and checkpoint sit in the same folder. In this repository, those files are organized under datasets/, config/, and model/, so update the file paths in the notebooks or set the working directory accordingly before running them locally.

Important Files

datasets/qoute_dataset.csv

Raw quote dataset used to build the next-word prediction samples.

config/X.npy

Preprocessed and padded input sequences with shape (85271, 745).

config/y.npy

Integer next-word labels with shape (85271,).

config/tokenizer.pkl

Saved tokenizer used for converting text to token ids and mapping predicted ids back to words.

config/config.pkl

Saved model and preprocessing constants.

model/next_word_checkpoint.pth

Saved PyTorch checkpoint containing the model state and optimizer state.

Notes

• main.py is currently a placeholder and does not run model inference.

• The dataset filename is spelled qoute_dataset.csv in the repository.

• The model is trained as a next-word classifier over a fixed vocabulary of 10,000 tokens.

• The current inference approach uses greedy decoding with argmax.

• The displayed training accuracy is training-set accuracy from the notebook loop, not a separate held-out test-set score.

• The X array ( config/X.npy ) is too big to upload to GitHub. So it is saved in the config folder in the compressed form (X.zip).