🚀 Next Word Prediction using LSTM (PyTorch)

A deep learning project that implements a Next Word Prediction model using LSTM (Long Short-Term Memory) networks in PyTorch. The model is trained on custom text data and learns to predict the most likely next word given a sequence of previous words.

---

📖 Project Overview

This project explores the fundamentals of language modeling and sequence learning using LSTMs.

Starting from raw text, the project:

• Tokenizes text using NLTK
• Builds a custom vocabulary
• Converts words into numerical representations
• Generates training sequences
• Pads variable-length sequences
• Trains an LSTM-based neural network
• Predicts the next word for a given input sequence

The project was built to gain a hands-on understanding of how language models work internally before moving on to more advanced architectures such as Transformers.

---

🛠️ Technologies Used

• Python
• PyTorch
• NLTK
• NumPy

---

🏗️ Model Architecture

Input Text
    │
    ▼
Tokenization
    │
    ▼
Vocabulary Encoding
    │
    ▼
Embedding Layer
(vocab_size → 100)
    │
    ▼
LSTM Layer
(input_size=100, hidden_size=150)
    │
    ▼
Fully Connected Layer
(150 → vocab_size)
    │
    ▼
Next Word Prediction

Model Configuration

Component	Value
Embedding Dimension	100
Hidden Size	150
Batch Size	32
Optimizer	Adam
Loss Function	CrossEntropyLoss
Framework	PyTorch

---

📂 Data Preparation

The training text is first tokenized and converted into numerical indices.

Example:

The train moved slowly

becomes:

[12, 45, 87, 103]

Training sequences are then generated incrementally:

[The, train]
[The, train, moved]
[The, train, moved, slowly]

The last word of each sequence is used as the target label.

---

🧠 Training Strategy

For each sequence:

Input  → All words except the last word
Target → Last word

Example:

Input:
The train moved

Target:
slowly

The model learns to predict the target word from the preceding context.

---

📈 Training

The model is trained using:

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

Training is performed on GPU when CUDA is available.

device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

---

🔮 Inference

Given a sequence of words:

The station was

The model outputs logits over the entire vocabulary and selects the word with the highest score:

value, index = torch.max(output, dim=1)

Example prediction:

The station was nearly

---

📚 Concepts Explored

This project helped me understand:

• Word Tokenization
• Vocabulary Construction
• Sequence Modeling
• Word Embeddings
• LSTM Internals
  • Hidden State
  • Cell State
  • Sequence Processing
• Cross Entropy Loss
• Batch Processing in PyTorch
• Dataset and DataLoader
• GPU Training with CUDA
• Neural Language Modeling

---

🚀 Future Improvements

Possible enhancements include:

• Multi-layer LSTM
• Bidirectional LSTM
• GRU-based Language Model
• Stateful LSTM
• Beam Search Decoding
• Training on larger datasets
• Transformer-based Next Word Prediction
• Interactive text generation interface

---

🎯 Example Workflow

Input Text
    │
    ▼
Tokenization
    │
    ▼
Vocabulary Creation
    │
    ▼
Sequence Generation
    │
    ▼
Padding
    │
    ▼
Embedding Layer
    │
    ▼
LSTM
    │
    ▼
Linear Layer
    │
    ▼
Vocabulary Scores
    │
    ▼
Predicted Next Word

---

👨‍💻 Author

Farhan Akbor Khan

Built as a learning project to gain practical experience with sequence models and neural language modeling using PyTorch.

---

⭐ If you found this project useful or interesting, consider starring the repository.