next-word-predictor-lstm

Next Word Predictor (LSTM)

🎯 A next word prediction model trained on WikiText-2 using LSTM architecture. Given a sequence of words, this model predicts the most likely next word.

📊 Model Performance

🚀 Quick Start

Installation

pip install torch huggingface_hub

Usage

import torch
import pickle
from huggingface_hub import hf_hub_download

# Download model files
model_path = hf_hub_download(
    repo_id="cosmicshubham/next-word-predictor-lstm",
    filename="best_model.pth"
)
tokenizer_path = hf_hub_download(
    repo_id="cosmicshubham/next-word-predictor-lstm",
    filename="tokenizer.pkl"
)

# Load tokenizer
with open(tokenizer_path, 'rb') as f:
    tokenizer = pickle.load(f)

# Define model architecture
import torch.nn as nn

class NextWordLSTM(nn.Module):
    def __init__(self, vocab_size, embedding_dim=256, hidden_dim=512, 
                 num_layers=2, dropout=0.5):
        super(NextWordLSTM, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.embedding_dropout = nn.Dropout(0.2)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers,
                           batch_first=True, 
                           dropout=dropout if num_layers > 1 else 0)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        
    def forward(self, x):
        embedded = self.embedding(x)
        embedded = self.embedding_dropout(embedded)
        lstm_out, _ = self.lstm(embedded)
        last_output = lstm_out[:, -1, :]
        dropped = self.dropout(last_output)
        output = self.fc(dropped)
        return output

# Load model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = NextWordLSTM(vocab_size=len(tokenizer.word2idx))

checkpoint = torch.load(model_path, map_location=device, weights_only=False)
model.load_state_dict(checkpoint['model_state_dict'])
model.to(device)
model.eval()

# Prediction function
def predict_next_word(text, top_k=5):
    SEQ_LENGTH = 20
    tokens = tokenizer.encode(text)
    
    if len(tokens) < SEQ_LENGTH:
        tokens = [0] * (SEQ_LENGTH - len(tokens)) + tokens
    else:
        tokens = tokens[-SEQ_LENGTH:]
    
    input_tensor = torch.tensor([tokens], dtype=torch.long).to(device)
    
    with torch.no_grad():
        output = model(input_tensor)
        probs = torch.softmax(output, dim=1)
        top_probs, top_indices = probs.topk(top_k * 3, dim=1)
    
    predictions = []
    for prob, idx in zip(top_probs[0], top_indices[0]):
        word = tokenizer.idx2word.get(idx.item(), '<UNK>')
        # Filter special tokens
        if word in ['<PAD>', '<UNK>', '<EOS>', '.', ',', '!', '?']:
            continue
        predictions.append((word, prob.item()))
        if len(predictions) >= top_k:
            break
    
    return predictions

# Example usage
text = "I am going to the"
predictions = predict_next_word(text, top_k=5)

print(f"Input: '{text}'\nPredictions:")
for i, (word, prob) in enumerate(predictions, 1):
    print(f"  {i}. {word:.<20} {prob*100:5.2f}%")

Expected Output

Input: 'I am going to the'
Predictions:
  1. right............... 0.93%
  2. city................ 0.85%
  3. north............... 0.74%
  4. end................. 0.71%
  5. new................. 0.70%

🏗️ Model Architecture

NextWordLSTM(
  (embedding): Embedding(10000, 256, padding_idx=0)
  (embedding_dropout): Dropout(p=0.2)
  (lstm): LSTM(256, 512, num_layers=2, batch_first=True, dropout=0.5)
  (dropout): Dropout(p=0.5)
  (fc): Linear(in_features=512, out_features=10000)
)

Total Parameters: ~13 Million

Architecture Details

• Embedding Dimension: 256
• Hidden Dimension: 512
• LSTM Layers: 2
• Dropout: 0.5 (LSTM) + 0.2 (Embedding)
• Vocabulary Size: 10,000 words
• Sequence Length: 20 tokens

📚 Training Details

Dataset

• Name: WikiText-2 (raw version)
• Source: Wikipedia articles
• Split: 70% Train, 15% Validation, 15% Test
• Training Samples: ~25,000 sequences
• Domain: Encyclopedic text

Hyperparameters

• Optimizer: Adam (lr=0.001, weight_decay=1e-5)
• Batch Size: 128
• Epochs: 20
• Scheduler: ReduceLROnPlateau (patience=2, factor=0.5)
• Early Stopping: Patience=5
• Gradient Clipping: Max norm=5.0

Regularization

• Dropout (0.5 on hidden layers, 0.2 on embeddings)
• L2 weight decay (1e-5)
• Gradient clipping
• Early stopping

📈 Training Results

The model was trained for 20 epochs with the following progression:

• Early stopping: No
• Best validation loss: 4.5036
• Final test perplexity: 88.31

🎯 Use Cases

1. Text Autocompletion: Suggest next words in writing applications
2. Search Query Completion: Predict search terms
3. Chatbot Enhancement: Improve response generation
4. Educational Tools: Language learning assistance
5. Text Generation: Create Wikipedia-style content

⚠️ Limitations

• Vocabulary: Limited to 10,000 most common words
• Context Window: Uses only last 20 tokens
• Domain: Optimized for Wikipedia-style formal text
• Out-of-vocabulary: Unknown words mapped to <UNK>
• Bias: Inherits biases from WikiText-2 dataset

🔬 Model Behavior

The model predicts words based on Wikipedia writing patterns:

• Strong at formal, encyclopedic language
• Predicts grammatically correct continuations
• Context-aware (e.g., location words after "going to the")
• Less effective with colloquial phrases or idioms

Example Predictions

Input: "I am going to the"

• ✅ "city", "north", "end" (contextually appropriate locations)

Input: "The weather today is"

• ✅ "expected", "forecast", "predicted" (formal weather descriptions)

📦 Files Included

• best_model.pth - Model checkpoint with best validation performance
• tokenizer.pkl - Vocabulary and tokenization mappings
• evaluation_metrics.json - Complete training metrics and history
• README.md - This documentation

🛠️ Training Code

The model was trained using custom LSTM implementation with:

• Sliding window input sequences
• Cross-entropy loss
• Top-k accuracy evaluation
• BLEU score computation

📊 Comparison with Baselines

📝 Citation

If you use this model in your research or application, please cite:

@misc{next-word-predictor-lstm-2026,
  author = {Shubham Kumar (cosmicshubham)},
  title = {Next Word Predictor: LSTM-based Language Model},
  year = {2025},
  month = {January},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/cosmicshubham/next-word-predictor-lstm}},
  note = {Trained on WikiText-2 dataset}
}

👤 Model Author

• Name: Shubham Kumar
• Username: cosmicshubham
• Contact: Hugging Face Profile

📄 License

MIT License - Free to use for commercial and non-commercial purposes.

🔗 Resources

• Model Repository
• WikiText-2 Dataset
• PyTorch Documentation

🙏 Acknowledgments

• WikiText-2 dataset creators
• PyTorch team
• Hugging Face for hosting infrastructure

---

Last Updated: 2026-01-24

Model Version: 1.0

Framework: PyTorch 1.9+