Building a Word-Level Language Model for Auto-Completion and Suggestions

This chapter covers the complete end-to-end architecture and implementation of a word-level language model engineered specifically for text auto-completion and predictive phrase suggestions. Using PyTorch, we build a multi-layer perceptron (MLP) language model from scratch, train it on over 73,000 real-world job titles, learn custom continuous word embeddings, and deploy a robust inference pipeline capable of real-time auto-completion.

1. Problem Statement and Task Formulation

Language models power core interactive digital workflows, ranging from IDE code completions and search engine queries to predictive mobile typing. At its core, an auto-completion engine takes a sequence of prefix words provided by a user and predicts the most statistically probable subsequent words based on patterns learned from a training corpus.

In contrast to character-level language models (which construct tokens character-by-character), a word-level language model operates directly on discrete word tokens. When trained on domain-specific datasets (such as professional job titles), the model learns domain semantics: given the input "software", it predicts job roles like "engineer" or "analyst" rather than arbitrary generic terms like "program" or "file".

Real-Life Analogy

Think of a word-level auto-completion engine like an administrative assistant handling phone routing. If you ask for the "Software" department, the assistant immediately offers options like "Software Engineering", "Software Quality Assurance", or "Software Architecture" because they have memorized the organization's exact operational structure.

2. Dataset Exploration and Lexical Processing

To train our model, we use the open-source BSTDS Job Titles Dataset hosted on Hugging Face. This dataset contains 73,380 standardized professional job title records.

Dataset Key Characteristics

• Total Records: 73,380 rows.
• Lexical Ordering: Sorted alphabetically/lexically (numbers first, followed by alphabetical characters).
• Sample Entries: "1st grade teacher", "1st pressman", "1st pressman web off- set".

Length Distribution Analysis

Understanding word count per sample is crucial for selecting an appropriate context window size.

# Calculating word distribution across job titles
word_count_map = {}
for title in job_titles:
    length = len(title.split())
    word_count_map[length] = word_count_map.get(length, 0) + 1

Key Observation: Approximately 80% to 90% of job titles consist of 2 or 3 words (e.g., "Software Engineer", "Web Solutions Developer"). This confirms that a compact sliding context window ($k = 3$) can capture the complete contextual history required to predict the next word.

3. Feature Engineering: Tokenization, Vocabulary, and Special Tokens

Neural networks cannot directly process raw strings; textual inputs must be converted into numerical index representations.

Step 1: Extracting Unique Vocabulary

We parse all 73,380 job titles, split them into constituent words, and compile a unique vocabulary set $V$:

Step 2: Special Start/Stop Token Insertion

To teach the model where a sequence begins and ends, we introduce a unified Start/Stop token represented by a period (.).

Note: The stop token (.) is critical for auto-completion systems. It teaches the network when a phrase is complete. For terminal words like "engineer" or "manager", the model learns to assign high probability to the stop token, signalling to the application that auto-completion should stop rather than suggest extraneous words.

The index for the special stop token is appended at index $\vert V_{raw}\vert = 12,157$, expanding our total vocabulary size $\vert V\vert$:

# Vocabulary mapping creation
unique_words = sorted(list(set(w for title in job_titles for w in title.split())))
vocab_size = len(unique_words) # 12,157

# Assign special stop token
STOP_TOKEN = "."
stop_token_idx = vocab_size # Index 12,157

# Bidirectional lookup dictionaries
word_to_idx = {w: i for i, w in enumerate(unique_words)}
word_to_idx[STOP_TOKEN] = stop_token_idx

idx_to_word = {i: w for i, w in enumerate(unique_words)}
idx_to_word[stop_token_idx] = STOP_TOKEN

4. Context Windowing and Sequence Tensor Construction

We structure the training data using a sliding context window of size $k = 3$. Each training sample consists of a fixed $k$-length input vector $X$ (the context) and a scalar target output $Y$ (the next word to predict).

Sliding Context Window Flow

For a sample title like "1st grade teacher" with stop token .:

# Context window tensor generation
block_size = 3  # Context window length (k = 3)
X, Y = [], []

for title in job_titles:
    words = title.split()
    context = [stop_token_idx] * block_size
    
    for w in words + [STOP_TOKEN]:
        target_idx = word_to_idx[w]
        X.append(context)
        Y.append(target_idx)
        # Shift context window to the left
        context = context[1:] + [target_idx]

import torch
X = torch.tensor(X, dtype=torch.long) # Shape: [269180, 3]
Y = torch.tensor(Y, dtype=torch.long) # Shape: [269180]

Dataset Volume: Processing 73,380 titles yields 269,180 training pairs.

5. Continuous Representation: Word Embedding Vectors

Mapping categorical word indices directly to integers (e.g., 'software' \to 11158) implies an artificial numerical ordering that neural networks cannot use effectively. Instead, we map each token index to a continuous $d$-dimensional embedding vector.

Embedding Lookup Mechanism

Let $d = 5$ be our embedding dimensionality. The lookup matrix $C$ holds trainable parameters:

For a context window of $k = 3$ words, each word index fetches a $5$-dimensional vector. Concatenating these vectors produces a single flattened input layer of dimension $k \times d = 3 \times 5 = 15$.

Embedding Hyperspace Intuition

During backpropagation, the model moves related words closer together in this 5D hyperspace. Words like "software", "web", and "systems" cluster near downstream targets like "engineer" or "developer", while remaining far from unrelated terms like "teacher" or "carpenter".

6. Model Architecture: 3-Layer Multi-Layer Perceptron (MLP)

We implement a deep neural network consisting of three fully connected hidden layers, each equipped with Batch Normalization (BatchNorm1d) and hyperbolic tangent activation functions (tanh), followed by a final linear multi-class output layer.

Mathematical Formulation per Layer

1. Flattened Embedding Input:

2. Hidden Layer 1:

3. Hidden Layer 2:

4. Hidden Layer 3:

5. Multi-Class Output Logits:

Model Architecture Diagram

import torch.nn as nn

class JobTitleLanguageModel(nn.Module):
    def __init__(self, vocab_size, emb_dim=5, block_size=3, hidden_dim=100):
        super().__init__()
        self.emb = nn.Embedding(vocab_size, emb_dim)
        
        input_dim = block_size * emb_dim # 3 * 5 = 15
        
        self.network = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.Tanh(),
            
            nn.Linear(hidden_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.Tanh(),
            
            nn.Linear(hidden_dim, hidden_dim),
            nn.BatchNorm1d(hidden_dim),
            nn.Tanh(),
            
            nn.Linear(hidden_dim, vocab_size) # Logits for all 12,158 words
        )

    def forward(self, x):
        # x shape: [B, 3]
        embeddings = self.emb(x) # Shape: [B, 3, 5]
        flattened = embeddings.view(embeddings.shape[0], -1) # Shape: [B, 15]
        logits = self.network(flattened) # Shape: [B, 12158]
        return logits

7. GPU Acceleration and Tensor Placement Workflows

Training on 269,180 sequences requires hardware acceleration. We use a cloud GPU environment (such as Kaggle's NVIDIA T4 allocation).

CUDA Setup and Device Allocation Workflow

To prevent CPU execution bottlenecks, all model parameters and dataset tensors must reside explicitly on the CUDA device.

# Hardware acceleration verification & tensor transfer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Executing pipeline on device: {device}")

# Instantiating model and pushing parameters to GPU
model = JobTitleLanguageModel(vocab_size=len(word_to_idx)).to(device)

# Moving dataset tensors to GPU memory
X = X.to(device)
Y = Y.to(device)

8. Model Training, Mini-Batching, and Loss Optimization

Loss Function

We use PyTorch's nn.CrossEntropyLoss(), which combines a log-softmax computation over the raw output logits $z$ with a negative log-likelihood loss against target indices $y$:

Mini-Batch Gradient Descent Strategy

• Batch Size ($B$): $512$ random samples per iteration.
• Learning Rate ($\eta$): $0.1$.
• Total Training Iterations: $100,000$.

import torch.optim as optimizer_lib

optimizer = optimizer_lib.SGD(model.parameters(), lr=0.1)
criterion = nn.CrossEntropyLoss()

loss_history = []
num_samples = X.shape[0]

for step in range(100000):
    # 1. Random mini-batch sampling
    ix = torch.randint(0, num_samples, (512,), device=device)
    X_batch, Y_batch = X[ix], Y[ix]
    
    # 2. Forward pass
    logits = model(X_batch)
    loss = criterion(logits, Y_batch)
    
    # 3. Backward pass & parameter optimization
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
    # Track metrics
    if step % 1000 == 0:
        loss_history.append(loss.item())
        print(f"Step {step:6d} | Loss: {loss.item():.4f}")

Training Dynamics and Loss Oscillation

Why does the loss plot fluctuate? Mini-batch sampling causes natural variance in the loss curve. Since each batch contains 512 randomly selected sequences, some batches will contain rare, difficult-to-predict terms while others contain common titles (e.g., "Software Engineer"). These random sample variations cause the observed oscillations, even as the overall loss trend decreases smoothly from $\approx 9.92$ down to $\approx 4.63$.

9. Auto-Completion Inference and Edge-Case Handling

Once trained, the model performs predictive generation via an autoregressive inference loop.

Autoregressive Generation Logic

1. Accept an initial seed word (e.g., "software").
2. Pad/construct an initial context vector of size $k = 3$.
3. Pass the context vector into the model to generate output logits over the vocabulary.
4. Select the most probable word index (using argmax or temperature sampling).
5. Append the predicted word to the context window and slide the window to the left.
6. Termination Conditions: Stop generating if:
   • The model predicts the special stop token (.), OR
   • The generated sequence reaches a maximum threshold length (e.g., $5$ words).

def generate_autocomplete(model, start_word, num_samples=10, max_words=5):
    model.eval()
    results = []
    
    with torch.no_grad():
        for _ in range(num_samples):
            # Check if start_word exists in vocab
            if start_word in word_to_idx:
                start_idx = word_to_idx[start_word]
                context = [stop_token_idx, stop_token_idx, start_idx]
                generated = [start_word]
            else:
                context = [stop_token_idx] * block_size
                generated = []

            for _ in range(max_words):
                x_input = torch.tensor([context], dtype=torch.long, device=device)
                logits = model(x_input)
                
                # Softmax probabilities & argmax selection
                probs = torch.softmax(logits, dim=1)
                predicted_idx = torch.argmax(probs, dim=1).item()
                
                # Check for termination token
                if predicted_idx == stop_token_idx:
                    break
                    
                predicted_word = idx_to_word[predicted_idx]
                generated.append(predicted_word)
                
                # Shift context window
                context = context[1:] + [predicted_idx]
                
            results.append(" ".join(generated))
            
    return results

Real-World Performance & Edge Cases

Case A: Seed Word = "software"

The model produces strong, highly relevant domain completions:

• "software engineer"
• "software analyst"
• "software developer"
• "software supervisor"
• "software technician"

(Note: Minor low-probability anomalies such as "software museum" can occasionally occur due to the compact MLP architecture).

Case B: Terminal Seed Words = "manager" or "engineer"

When given words that typically end a job title, the model correctly predicts the stop token (.) immediately, returning no extra completion terms:

• Input: "manager" $\rightarrow$ Output: "manager" (0 extra words added)
• Input: "engineer" $\rightarrow$ Output: "engineer" (0 extra words added)

Case C: Unconstrained Sampling (Start Token .)

When seeded with only the stop token, the model generates diverse full titles from scratch:

• "carpenter"
• "prepress technician"
• "warehouse associate"
• "dean"

10. Architectural Evolution: From MLP to Transformers and LLMs

The model built in this chapter illustrates the fundamental mechanics of statistical auto-completion. Below is a structural comparison showing how this lightweight MLP architecture compares to modern production-grade models.

To scale our lightweight engine into a modern Large Language Model (LLM):

1. Expand Context and Data: Train on billions of web-scale text tokens rather than narrow job titles.
2. Upgrade Backbone: Replace static concatenated MLPs with Transformer Self-Attention layers.
3. Instruction Fine-Tuning: Apply Reinforcement Learning from Human Feedback (RLHF) to transition the engine from basic text completion to an interactive conversational AI.

Summary and Key Takeaways

• Word-Level vs. Character-Level Modeling: Word-level language models operate directly on discrete token dictionaries ($\vert V\vert \approx 12,158$), making them faster and better suited for word-level auto-completion than character-level models.
• The Role of Special Stop Tokens: Unified start/stop tokens (.) are essential. They teach the network when a phrase is complete, preventing infinite generation loops on naturally terminal words like "engineer" or "manager".
• Dense Embeddings vs. Raw Indices: Raw numerical IDs carry no semantic distance. Mapping word indices into continuous $d$-dimensional vector spaces ($d=5$) allows backpropagation to cluster syntactically and semantically related words together.
• Context Windowing: Sliding a fixed context window ($k=3$) across token sequences turns unstructured text corpora into structured supervised training pairs ($X \in \mathbb{R}^{B \times 3}, Y \in \mathbb{R}^{B}$).
• 3-Layer MLP Design: Combining Fully Connected layers with Batch Normalization (BatchNorm1d) and non-linear activations (tanh) stabilizes gradient flow and speeds up convergence across mini-batches.
• GPU Memory Transfers: To avoid CPU bottlenecks, both the model parameters (model.to('cuda')) and the dataset tensors (X.to('cuda')) must be placed explicitly on the GPU device.
• Mini-Batching Dynamics: Optimization using mini-batches ($B=512$) causes expected variance and fluctuations in the training loss curve while maintaining an overall downward trend.
• Inference Autoregression: Predictive auto-completion works by feeding predicted tokens back into the sliding context window until a stop token is hit or a maximum word limit is reached.