app.py

"""Next word prediction with trigrams counted on the Selma Lagerlof corpus.

EDAN20, lab 2.
"""

import json

import gradio as gr
from huggingface_hub import hf_hub_download

# Change this to your own dataset repository
DATASET = "YOUR_USERNAME/selma_ngrams"
CAND_NBR = 5


def load_ngrams_jsonl(file_name):
    dictionary = {}
    with open(file_name, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            data = json.loads(line)
            k = tuple(data["ngram"]) if len(data["ngram"]) > 1 else data["ngram"][0]
            dictionary[k] = data["count"]
    return dictionary


frequency = load_ngrams_jsonl(
    hf_hub_download(DATASET, "unigrams.jsonl", repo_type="dataset"))
frequency_bigrams = load_ngrams_jsonl(
    hf_hub_download(DATASET, "bigrams.jsonl", repo_type="dataset"))
frequency_trigrams = load_ngrams_jsonl(
    hf_hub_download(DATASET, "trigrams.jsonl", repo_type="dataset"))


def predict(text):
    """Predict the five most likely next words with a trigram model."""
    tokens = text.lower().split()
    # We need two words of context; <s> stands for the start of the sentence
    context = tuple((['<s>', '<s>'] + tokens)[-2:])
    candidates = {trigram: count for trigram, count in frequency_trigrams.items()
                  if trigram[:2] == context}
    if not candidates:
        # No trigram seen with this context, back off to the bigrams
        candidates = {bigram: count for bigram, count in frequency_bigrams.items()
                      if bigram[0] == context[1]}
        sorted_candidates = sorted(candidates.items(), key=lambda x: (-x[1], x[0][1]))
        best = [(ngram[1], count) for ngram, count in sorted_candidates[:CAND_NBR]]
    else:
        sorted_candidates = sorted(candidates.items(), key=lambda x: (-x[1], x[0][2]))
        best = [(ngram[2], count) for ngram, count in sorted_candidates[:CAND_NBR]]

    if not best:
        return "No prediction, this context is not in the corpus."
    return "\n".join("{}\t{}".format(word, count) for word, count in best)


demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Beginning of a sentence", value="Det var en"),
    outputs=gr.Textbox(label="Next word candidates and their counts"),
    title="Selma Lagerlof trigram language model",
    description="Type the beginning of a Swedish sentence and the model "
                "proposes the five most frequent continuations.")

if __name__ == "__main__":
    demo.launch()