Next Word Predictor

easyFREELanguage ModelingNLP 3: Advanced

Given a corpus and a context word, predict the most likely next word using bigram counts.

If there are ties, return the word that comes first alphabetically.

Input format:

Line 1: The training corpus (lowercase)
Line 2: The context word

Output: The predicted next word and its probability (rounded to 4 decimal places but not padded with zeros: $2/3$ prints as 0.6667, $1/2$ as 0.5, and $1$ as 1.0), separated by a space.

Example

Input

i like cats i like dogs i hate rain
i

Output

like 0.6667

Reasoning

Inputs: the corpus is the string i like cats i like dogs i hate rain, nine words in order, and the context is i. A bigram is a pair of adjacent words $(w_k, w_{k+1})$ , so this corpus contains 8 of them.
Bigrams that start with the context: i appears three times, and the word right after it is "like", "like", "hate". So $C(i, \text{like}) = 2$ , $C(i, \text{hate}) = 1$ , and the total is $C(i) = 3$ .
Prediction: the most frequent follower is "like" (2 beats 1), so no tie-break is needed here. If "hate" also had count 2, alphabetical order would pick "hate".
Probability: $P(\text{like} \mid i) = \frac{C(i, \text{like})}{C(i)} = \frac{2}{3} \approx 0.6667$ after rounding to 4 decimal places, giving like 0.6667.
Notice: the denominator counts only occurrences of the context that are followed by some word. A context that sits only at the very end of the corpus has no followers at all.

Constraints

All text is lowercase
Return most frequent next word after the context word
Ties broken alphabetically
Output: "word probability"
If context word never appears, output "N/A 0.0"

Learn this in

Introduction to LLMsLLM Study Plan Text RepresentationNLP Study Plan Language ModelingNLP Study Plan

More from NLP 3: Advanced

Temperature Samplereasy ROUGE-1 Score Calculatoreasy Context Window Findereasy Keyword Answer Extractoreasy Levenshtein Distanceeasy

Concept primer

Background Knowledge

Language Modeling is a fundamental task in Natural Language Processing (NLP) that involves predicting the next word in a sequence of words. This task is crucial for various applications, such as text generation, language translation, and speech recognition. A key concept in language modeling is the use of n-grams, which are sequences of $n$ items from a given text. In this problem, we will be working with bigrams, which are sequences of 2 items.

The probability of a word given its context is calculated using the conditional probability formula: $P(w_2|w_1) = \frac{P(w_1, w_2)}{P(w_1)}$ , where $w_1$ is the context word and $w_2$ is the word we want to predict. To calculate these probabilities, we need to count the occurrences of bigrams in the training corpus. This is where bigram counts come into play. By counting the occurrences of each bigram, we can estimate the probability of a word given its context.

In this problem, we will use the Maximum Likelihood Estimation (MLE) approach to predict the next word. The MLE approach involves finding the word that maximizes the probability $P(w_2|w_1)$ . If there are ties, we will return the word that comes first alphabetically. This approach is simple and effective, but it has its limitations. For example, it does not take into account the context beyond the immediate previous word.

Algorithm/Approach

The general approach to solve this problem involves the following steps:

Preprocess the training corpus to extract bigrams
Calculate the bigram counts and probabilities
Use the MLE approach to predict the next word given the context word
Handle ties by returning the word that comes first alphabetically

This approach is a simple example of a language model, which is a probabilistic model that predicts the next word in a sequence of words.

Step-by-Step Strategy

To implement the solution, follow these steps:

Preprocess the training corpus: Convert the corpus to lowercase and split it into individual words.
Extract bigrams: Iterate through the list of words and extract bigrams.
Calculate bigram counts: Count the occurrences of each bigram.
Calculate probabilities: Calculate the probability of each word given the context word using the bigram counts.
Predict the next word: Use the MLE approach to predict the next word given the context word.
Handle ties: If there are ties, return the word that comes first alphabetically.

Common Pitfalls

Some common pitfalls to watch out for when implementing this solution include:

Not handling out-of-vocabulary words
Not handling ties correctly
Not normalizing the probabilities correctly
Not using the correct data structures to store the bigram counts and probabilities

Time & Space Complexity

The time complexity of this solution is $O(n)$ , where $n$ is the length of the training corpus. This is because we need to iterate through the corpus to extract bigrams and calculate bigram counts. The space complexity is also $O(n)$ , as we need to store the bigram counts and probabilities in memory. However, the space complexity can be reduced by using a more efficient data structure, such as a hash table, to store the bigram counts and probabilities.

solution.py

Python 3.13

Test results