Detect the language of a text

"""
Detect the language of a text: character trigram profiles, rank distance.

Rung N0. Deterministic, standard library only, and the whole model is a few
hundred short strings per language.

The idea is older than most of the field. Every language repeats its own
trigrams: « ent », « les », « eur » in French, « the », « ing » in English,
« que », « los » in Spanish. Rank those trigrams by frequency in a sample of
the language, rank them again in the text to identify, and compare the two
orderings. The language whose ordering is closest wins.

Two details make it work.

First, words are padded with spaces before being cut, so a trigram carries
the information that it opens or closes a word. « les » inside a word is not
the article.

Second, the comparison is on ranks, not on frequencies. A rank survives a
sample four times longer, and a text four times shorter, unchanged.
"""

import re
import unicodedata
from collections import Counter

PROFILE_SIZE = 300

# Letters only. Digits, punctuation and symbols say nothing about a language,
# and a text full of them would drown the trigrams that do.
WORDS = re.compile(r"[^\W\d_]+")


def trigrams(text: str):
    """Yield the padded trigrams of every word, in reading order."""
    lowered = unicodedata.normalize("NFC", text.lower())
    for word in WORDS.findall(lowered):
        padded = f" {word} "
        for i in range(len(padded) - 2):
            yield padded[i:i + 3]


def profile(sample: str, size: int = PROFILE_SIZE) -> dict[str, int]:
    """
    Build a language profile from a sample: trigram to rank, most frequent
    first.

    Ties are broken alphabetically, so the same sample always gives the same
    profile. A language model you cannot reproduce is a language model you
    cannot debug.
    """
    counts = Counter(trigrams(sample))
    ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
    return {gram: rank for rank, (gram, _) in enumerate(ordered[:size])}


def distance(text: str, reference: dict[str, int], size: int = PROFILE_SIZE) -> float:
    """
    Out-of-place distance between the text and one language profile.

    Every trigram of the text costs how far it moved in the ranking. A
    trigram the language never uses costs the maximum, which is what makes an
    unrelated language expensive rather than merely different.

    Dividing by the number of trigrams keeps a long text and a short one on
    the same scale.
    """
    text_profile = profile(text, size)
    if not text_profile:
        return float(size)
    total = 0
    for gram, rank in text_profile.items():
        reference_rank = reference.get(gram)
        total += size if reference_rank is None else abs(rank - reference_rank)
    return total / len(text_profile)


def ranked(text: str, profiles: dict[str, dict[str, int]]) -> list[tuple[str, float]]:
    """
    Every candidate language, closest first.

    The caller gets the gap between the first two, which is the only honest
    measure of how sure this is.
    """
    scores = [(name, distance(text, reference)) for name, reference in profiles.items()]
    return sorted(scores, key=lambda item: (item[1], item[0]))


def detect(text: str, profiles: dict[str, dict[str, int]]) -> str:
    """The closest language. It always returns one, even when it should not."""
    return ranked(text, profiles)[0][0]

/**
 * Detect the language of a text: character trigram profiles, rank distance.
 *
 * Rung N0. Deterministic, no dependency, and the whole model is a few hundred
 * short strings per language.
 *
 * The idea is older than most of the field. Every language repeats its own
 * trigrams: "ent", "les", "eur" in French, "the", "ing" in English, "que",
 * "los" in Spanish. Rank those trigrams by frequency in a sample of the
 * language, rank them again in the text to identify, and compare the two
 * orderings. The language whose ordering is closest wins.
 *
 * Two details make it work.
 *
 * First, words are padded with spaces before being cut, so a trigram carries
 * the information that it opens or closes a word. "les" inside a word is not
 * the article.
 *
 * Second, the comparison is on ranks, not on frequencies. A rank survives a
 * sample four times longer, and a text four times shorter, unchanged.
 */

export const PROFILE_SIZE = 300;

// Letters only. Digits, punctuation and symbols say nothing about a language,
// and a text full of them would drown the trigrams that do.
const WORDS = /\p{L}+/gu;

/** The padded trigrams of every word, in reading order. */
export function trigrams(text) {
  const grams = [];
  for (const word of text.normalize('NFC').toLowerCase().match(WORDS) ?? []) {
    const padded = ` ${word} `;
    for (let i = 0; i + 3 <= padded.length; i += 1) grams.push(padded.slice(i, i + 3));
  }
  return grams;
}

/**
 * Build a language profile from a sample: trigram to rank, most frequent
 * first.
 *
 * Ties are broken alphabetically, so the same sample always gives the same
 * profile. A language model you cannot reproduce is a language model you
 * cannot debug.
 */
export function profile(sample, size = PROFILE_SIZE) {
  const counts = new Map();
  for (const gram of trigrams(sample)) counts.set(gram, (counts.get(gram) ?? 0) + 1);
  const ordered = [...counts].sort((a, b) => b[1] - a[1] || (a[0] < b[0] ? -1 : 1));
  return new Map(ordered.slice(0, size).map(([gram], rank) => [gram, rank]));
}

/**
 * Out-of-place distance between the text and one language profile.
 *
 * Every trigram of the text costs how far it moved in the ranking. A trigram
 * the language never uses costs the maximum, which is what makes an unrelated
 * language expensive rather than merely different.
 *
 * Dividing by the number of trigrams keeps a long text and a short one on the
 * same scale.
 */
export function distance(text, reference, size = PROFILE_SIZE) {
  const textProfile = profile(text, size);
  if (textProfile.size === 0) return size;
  let total = 0;
  for (const [gram, rank] of textProfile) {
    const referenceRank = reference.get(gram);
    total += referenceRank === undefined ? size : Math.abs(rank - referenceRank);
  }
  return total / textProfile.size;
}

/**
 * Every candidate language, closest first.
 *
 * The caller gets the gap between the first two, which is the only honest
 * measure of how sure this is.
 *
 * @param {string} text
 * @param {Map<string, Map<string, number>>} profiles
 */
export function ranked(text, profiles) {
  const scores = [...profiles].map(([name, reference]) => [name, distance(text, reference)]);
  return scores.sort((a, b) => a[1] - b[1] || (a[0] < b[0] ? -1 : 1));
}

/** The closest language. It always returns one, even when it should not. */
export function detect(text, profiles) {
  return ranked(text, profiles)[0][0];
}

"""
Detect the language of a text with a naive Bayes classifier on character
n-grams.

Rung N1. Same features as N0, one to three characters, but weighed instead of
ranked. Each n-gram of the text votes for every language, in proportion to
how often that language uses it, and the votes are multiplied together.

What this buys over the rank distance of N0 is a number the caller can act
on. N0 answers « French »; this answers « French, and here is how far ahead
of Spanish it is ». A detector that can abstain is worth more than one that
is right slightly more often.

Training is a paragraph per language and a fraction of a second. The model is
a table of counts.
"""

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline


def train(samples: dict[str, str]):
    """
    Fit on one sample of text per language.

    `char_wb` cuts n-grams inside word boundaries, so an n-gram carries the
    information that it opens or closes a word, exactly as the padding of N0
    did. Smoothing is light: an n-gram this language never used should count
    against it, without ruling it out on a single character.
    """
    model = make_pipeline(
        CountVectorizer(analyzer="char_wb", ngram_range=(1, 3), lowercase=True),
        MultinomialNB(alpha=0.1),
    )
    model.fit(list(samples.values()), list(samples.keys()))
    return model


def probabilities(model, text: str) -> dict[str, float]:
    """
    How the text splits between the known languages.

    Read these as an ordering, not as a measure of truth: they always sum to
    one, over the languages the model was trained on and no others.
    """
    scores = model.predict_proba([text])[0]
    return {str(name): float(score) for name, score in zip(model.classes_, scores)}


def detect(model, text: str, minimum: float = 0.0) -> str | None:
    """
    The most likely language, or None when the model is not sure enough.

    `minimum` is yours to set. Raise it when a wrong language costs more than
    no answer, for instance when the answer picks the queue a message is
    routed to. Leave it at zero to always get a name, as N0 does.
    """
    best, score = max(probabilities(model, text).items(), key=lambda item: item[1])
    return best if score >= minimum else None

/**
 * Detect the language of a text with a naive Bayes classifier on character
 * n-grams.
 *
 * Rung N1. Same features as N0, one to three characters, but weighed instead
 * of ranked. Each n-gram of the text votes for every language, in proportion
 * to how often that language uses it, and the votes are multiplied together.
 *
 * What this buys over the rank distance of N0 is a number the caller can act
 * on. N0 answers "French"; this answers "French, and here is how far ahead of
 * Spanish it is". A detector that can abstain is worth more than one that is
 * right slightly more often.
 *
 * Written out in full rather than pulled from a library, because multinomial
 * naive Bayes is thirty lines of counting. That is the argument of this rung:
 * the classical tool is small enough to read.
 */

// Smoothing. An n-gram a language never used should count against it, without
// ruling the language out on the strength of a single character.
const ALPHA = 0.1;

/**
 * The n-grams of one to three characters of every word, padded with spaces so
 * that an n-gram carries the information that it opens or closes a word.
 */
export function ngrams(text) {
  const grams = [];
  for (const word of text.toLowerCase().split(/\s+/).filter(Boolean)) {
    const padded = ` ${word} `;
    for (let n = 1; n <= 3; n += 1) {
      for (let i = 0; i + n <= padded.length; i += 1) grams.push(padded.slice(i, i + n));
    }
  }
  return grams;
}

/** Fit on one sample of text per language: `{ fr: "…", en: "…" }`. */
export function train(samples) {
  const vocabulary = new Set();
  const counts = new Map();
  for (const [language, sample] of Object.entries(samples)) {
    const perLanguage = new Map();
    for (const gram of ngrams(sample)) {
      perLanguage.set(gram, (perLanguage.get(gram) ?? 0) + 1);
      vocabulary.add(gram);
    }
    counts.set(language, perLanguage);
  }
  return { counts, vocabulary };
}

/**
 * How the text splits between the known languages.
 *
 * Read these as an ordering, not as a measure of truth: they always sum to
 * one, over the languages the model was trained on and no others.
 */
export function probabilities(model, text) {
  const { counts, vocabulary } = model;
  const grams = ngrams(text).filter((gram) => vocabulary.has(gram));
  const logScores = [...counts].map(([language, perLanguage]) => {
    const total = [...perLanguage.values()].reduce((a, b) => a + b, 0);
    const denominator = total + ALPHA * vocabulary.size;
    // Sum of logs rather than a product of probabilities: multiplying a few
    // thousand small numbers underflows to zero.
    let score = 0;
    for (const gram of grams) score += Math.log(((perLanguage.get(gram) ?? 0) + ALPHA) / denominator);
    return [language, score];
  });
  const highest = Math.max(...logScores.map(([, score]) => score));
  const weights = logScores.map(([language, score]) => [language, Math.exp(score - highest)]);
  const sum = weights.reduce((a, [, weight]) => a + weight, 0);
  return Object.fromEntries(weights.map(([language, weight]) => [language, weight / sum]));
}

/**
 * The most likely language, or null when the model is not sure enough.
 *
 * `minimum` is yours to set. Raise it when a wrong language costs more than no
 * answer, for instance when the answer picks the queue a message is routed to.
 * Leave it at zero to always get a name, as N0 does.
 */
export function detect(model, text, minimum = 0) {
  const scores = Object.entries(probabilities(model, text));
  const [best, score] = scores.sort((a, b) => b[1] - a[1] || (a[0] < b[0] ? -1 : 1))[0];
  return score >= minimum ? best : null;
}

"""
Detect the language of a text by asking a general-purpose model.

Rung N3. This is the option people reach for first. It is here so you can see
what it costs, not because this entry recommends it.

Note what the code has to do that N0 did not: cap the input, send only an
excerpt, retry on failure, parse an answer that is only probably valid JSON,
normalise a code the model may write in half a dozen ways, and refuse an
answer that is outside the list it was given. That plumbing is the real cost
of this rung, and it is the part your tests have to cover, because the model
itself is not testable.

The one thing this rung genuinely adds is that it needs no sample of the
language. The one thing it cannot do is tell you it is wrong.
"""

from __future__ import annotations

import json
from collections.abc import Collection

PROMPT = (
    "Identify the language of the text below.\n"
    "Answer with JSON only: an object with keys `language` and `confidence`,\n"
    "where `language` is a two-letter ISO 639-1 code chosen from this list:\n"
    "{languages}, or `und` if the text is in none of them.\n\n"
    "Text:\n{excerpt}"
)

MAX_CHARACTERS = 8000

# A language is decided in the first few sentences. Sending the whole document
# is not thoroughness, it is paying by the token for nothing.
EXCERPT_CHARACTERS = 600


class DetectionUnavailable(Exception):
    """The provider could not be reached, or answered something unusable."""


def detect(text: str, languages: Collection[str], client=None, *, attempts: int = 3) -> str | None:
    """
    Return the code of the detected language, or None when the model says the
    text is in none of the languages it was offered.

    `client` is injected so this function can be tested without a network
    call. In production it defaults to a real provider client.
    """
    if client is None:  # pragma: no cover - needs a key and a network
        from openai import OpenAI

        client = OpenAI()

    # A model charges by the token. Refusing oversized input is not an
    # optimisation, it is a cost control.
    if len(text) > MAX_CHARACTERS:
        raise ValueError(f"text longer than {MAX_CHARACTERS} characters")

    answer = _ask(client, text[:EXCERPT_CHARACTERS], sorted(languages), attempts)

    # Models answer « fr », « FR », « fr-CA » and « French » for the same
    # thing. Everything but the first is a bug waiting to reach production.
    code = str(answer.get("language", "")).strip().lower().split("-")[0]
    if code == "und":
        return None
    if code not in languages:
        raise DetectionUnavailable(f"the model answered a language outside the list: {code!r}")
    return code


def _ask(client, excerpt: str, languages: list[str], attempts: int) -> dict:
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            answer = client.complete(
                prompt=PROMPT.format(languages=", ".join(languages), excerpt=excerpt),
                # Temperature zero, because a routing decision that changes
                # between two identical calls cannot be reviewed.
                temperature=0,
            )
            parsed = json.loads(answer)
            if isinstance(parsed, dict):
                return parsed
            last_error = ValueError("the model answered something that is not an object")
        except Exception as error:  # noqa: BLE001 - any provider failure is retried
            last_error = error
    raise DetectionUnavailable(str(last_error))

/**
 * Detect the language of a text by asking a general-purpose model.
 *
 * Rung N3. This is the option people reach for first. It is here so you can
 * see what it costs, not because this entry recommends it.
 *
 * Note what the code has to do that N0 did not: cap the input, send only an
 * excerpt, retry on failure, parse an answer that is only probably valid
 * JSON, normalise a code the model may write in half a dozen ways, and refuse
 * an answer that is outside the list it was given. That plumbing is the real
 * cost of this rung, and it is the part your tests have to cover, because the
 * model itself is not testable.
 *
 * The one thing this rung genuinely adds is that it needs no sample of the
 * language. The one thing it cannot do is tell you it is wrong.
 */

export const MAX_CHARACTERS = 8000;

// A language is decided in the first few sentences. Sending the whole
// document is not thoroughness, it is paying by the token for nothing.
export const EXCERPT_CHARACTERS = 600;

export class DetectionUnavailable extends Error {}

/** The exact request sent to the model. Exported so a test can read it. */
export function buildPrompt(languages, excerpt) {
  return [
    'Identify the language of the text below.',
    'Answer with JSON only: an object with keys `language` and `confidence`,',
    'where `language` is a two-letter ISO 639-1 code chosen from this list:',
    `${[...languages].sort().join(', ')}, or \`und\` if the text is in none of them.`,
    '',
    'Text:',
    excerpt,
  ].join('\n');
}

/**
 * The code of the detected language, or null when the model says the text is
 * in none of the languages it was offered.
 *
 * @param {string} text
 * @param {string[]} languages the codes the model may choose from
 * @param {object} options
 * @param {{complete: Function}} [options.client] injected so this can be
 *   tested without a network call; defaults to a real provider client
 * @param {number} [options.attempts]
 */
export async function detect(text, languages, { client, attempts = 3 } = {}) {
  if (!client) {
    // Needs a key and a network, so it is never reached in the tests.
    const { OpenAI } = await import('openai');
    client = new OpenAI();
  }

  // A model charges by the token. Refusing oversized input is not an
  // optimisation, it is a cost control.
  if (text.length > MAX_CHARACTERS) {
    throw new RangeError(`text longer than ${MAX_CHARACTERS} characters`);
  }

  const answer = await ask(client, text.slice(0, EXCERPT_CHARACTERS), languages, attempts);

  // Models answer "fr", "FR", "fr-CA" and "French" for the same thing.
  // Everything but the first is a bug waiting to reach production.
  const code = String(answer.language ?? '').trim().toLowerCase().split('-')[0];
  if (code === 'und') return null;
  if (![...languages].includes(code)) {
    throw new DetectionUnavailable(`the model answered a language outside the list: ${code}`);
  }
  return code;
}

async function ask(client, excerpt, languages, attempts) {
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      const answer = await client.complete({
        prompt: buildPrompt(languages, excerpt),
        // Temperature zero, because a routing decision that changes between
        // two identical calls cannot be reviewed.
        temperature: 0,
      });
      const parsed = JSON.parse(answer);
      if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) return parsed;
      lastError = new Error('the model answered something that is not an object');
    } catch (error) {
      lastError = error;
    }
  }
  throw new DetectionUnavailable(String(lastError));
}