This repository contains an LSTM-based next word prediction model implemented in PyTorch.
The model utilizes advanced techniques including an extra fully connected layer with ReLU and dropout, layer normalization, label smoothing loss, gradient clipping, and learning rate scheduling to improve performance.
It also uses SentencePiece for subword tokenization.
Model Details
Model Description
The LSTM Next Word Predictor is designed to predict the next word or subword given an input sentence.
The model is trained on a dataset provided in CSV format (with a 'data' column) and uses an LSTM network with many enhancements.
This model can be directly used for next word prediction in text autocompletion.
Downstream Use
The model can be fine-tuned for related tasks such as:
Text generation.
Language modeling for specific domains.
Out-of-Scope Use
This model is not suitable for:
Tasks requiring deep contextual understanding beyond next-word prediction.
Applications where transformer-based architectures are preferred for longer contexts.
Sensitive applications where data bias could lead to unintended outputs.
Risks and Limitations
Risks: Inaccurate or unexpected predictions may occur if the input context is too complex or ambiguous.
Limitations: The model’s performance is bounded by the size and quality of the training data as well as the inherent limitations of LSTM architectures in modeling long-range dependencies.
Recommendations
Users should be aware of the above limitations and conduct appropriate evaluations before deploying the model in production.
Consider further fine-tuning or additional data preprocessing if the model is applied in sensitive contexts.
How to Get Started with the Model
Use the code below to get started with the model.
To run the model, follow these steps:
Training:
Ensure you have a CSV file with a column named data containing your training sentences.
The model will output the top predicted word or subword.
Training Details
Training Data
Data Source: CSV file with a column data containing sentences.
Preprocessing: Uses SentencePiece for subword tokenization.
Dataset: The training and validation datasets are split based on a user-defined ratio.
Training Procedure
Preprocessing: Tokenization using a SentencePiece model.
Training Hyperparameters:
Batch Size: Configurable via --batch_size (default: 512)
Learning Rate: Configurable via --learning_rate (default: 0.001)
Epochs: Configurable via --num_epochs (default: 25)
LSTM Parameters: Configurable number of layers, dropout, and hidden dimensions.
Label Smoothing: Applied with a configurable factor (default: 0.1)
Optimization: Uses Adam optimizer with weight decay and gradient clipping.
Learning Rate Scheduling: ReduceLROnPlateau scheduler based on validation loss.
Speeds, Sizes, Times
Checkpoint and TorchScript models are saved during training for later inference.
Evaluation
Testing Data, Factors & Metrics
Testing Data: Derived from the same CSV, split from the training data.
Metrics: Primary metric is the loss (with label smoothing), with qualitative evaluation based on next-word accuracy.
Factors: Evaluations may vary based on sentence length and dataset diversity.
Summary
The model demonstrates promising performance on next word prediction tasks;
however, quantitative results (e.g., accuracy, loss) should be validated on your specific dataset.
Model Examination
Interpretability techniques such as examining predicted token distributions can be applied to further understand model behavior.
Technical Specifications
Model Architecture and Objective
Architecture: LSTM-based network with enhancements such as an extra fully connected layer, dropout, and layer normalization.
Objective: Predict the next word/subword given a sequence of tokens.
Citation
BibTeX:
bibtex
1@misc {aarohan_verma_2025,
2 author = { {Aarohan Verma} },
3 title = { lstm-next-word-predictor },
4 year = 2025,
5 url = { https://huggingface.co/aarohanverma/lstm-next-word-predictor },
6 doi = { 10.57967/hf/4882 },
7 publisher = { Hugging Face }
8}
Model Card Contact
For inquiries or further information, please contact:
