Tag articles by topic

"""
Tag articles from a controlled vocabulary, matched after simple lemmatisation.

Rung N0. Deterministic, standard library only, and auditable: every tag can be
traced back to the term that produced it, which is what an editor asks for the
first time a tag looks wrong.

Two things make this work.

First, the article and the vocabulary go through the very same pipeline. A
stemmer that is applied to one side only will happily fail to match a word
with itself.

Second, tagging is multi-label by construction. An article covers several
topics, or none; a function that returns one topic per article is answering a
question nobody asked.
"""

import re
import unicodedata

# Endings stripped, longest first. This is a plural-and-suffix stripper, not a
# linguist's lemmatiser. It only has to fold the spellings of one word onto
# each other, and both sides of the comparison get the same treatment.
SUFFIXES = ("ements", "ement", "ations", "ation", "es", "s", "x", "e")

WORD = re.compile(r"[^\W_]+")


def normalise(text: str) -> str:
    """Lowercase and drop accents, so « Fiscalité » and « FISCALITE » meet."""
    decomposed = unicodedata.normalize("NFD", text.lower())
    return "".join(c for c in decomposed if not unicodedata.combining(c))


def lemmatise(word: str) -> str:
    """Strip one ending, and only when a stem of three letters is left."""
    for suffix in SUFFIXES:
        if word.endswith(suffix) and len(word) - len(suffix) >= 3:
            return word[: -len(suffix)]
    return word


def stems(text: str) -> str:
    """
    The text as a run of stems, padded with spaces at both ends.

    The padding is what lets a multi-word term be found with a plain substring
    search: « impot » can then never match inside « impotent ».
    """
    return " " + " ".join(lemmatise(w) for w in WORD.findall(normalise(text))) + " "


def tag(article: str, vocabulary: dict[str, list[str]], min_terms: int = 1) -> list[str]:
    """
    Every topic whose terms appear in the article, best supported first.

    `vocabulary` maps a topic name to the terms that stand for it, single or
    multi-word. Keeping it a parameter is the point of this rung: the taxonomy
    belongs to whoever edits the articles, not to the code.

    `min_terms` is how many distinct terms a topic needs before it is claimed.
    Raise it when a single passing mention is not enough to file an article.
    """
    haystack = stems(article)
    hits = {}
    for topic, terms in vocabulary.items():
        found = sum(1 for term in terms if stems(term) in haystack)
        if found >= min_terms:
            hits[topic] = found
    # Best supported first, ties in alphabetical order: a tagging run has to be
    # replayable, and a set has no order to replay.
    return sorted(hits, key=lambda topic: (-hits[topic], topic))

/**
 * Tag articles from a controlled vocabulary, matched after simple lemmatisation.
 *
 * Rung N0. Deterministic, no dependency, and auditable: every tag can be
 * traced back to the term that produced it, which is what an editor asks for
 * the first time a tag looks wrong.
 *
 * Two things make this work.
 *
 * First, the article and the vocabulary go through the very same pipeline. A
 * stemmer that is applied to one side only will happily fail to match a word
 * with itself.
 *
 * Second, tagging is multi-label by construction. An article covers several
 * topics, or none; a function that returns one topic per article is answering
 * a question nobody asked.
 */

// Endings stripped, longest first. This is a plural-and-suffix stripper, not a
// linguist's lemmatiser. It only has to fold the spellings of one word onto
// each other, and both sides of the comparison get the same treatment.
const SUFFIXES = ['ements', 'ement', 'ations', 'ation', 'es', 's', 'x', 'e'];

const WORD = /[\p{L}\p{N}]+/gu;

/** Lowercase and drop accents, so « Fiscalité » and « FISCALITE » meet. */
export function normalise(text) {
  return text.toLowerCase().normalize('NFD').replace(/\p{M}+/gu, '');
}

/** Strip one ending, and only when a stem of three letters is left. */
export function lemmatise(word) {
  for (const suffix of SUFFIXES) {
    if (word.endsWith(suffix) && word.length - suffix.length >= 3) {
      return word.slice(0, -suffix.length);
    }
  }
  return word;
}

/**
 * The text as a run of stems, padded with spaces at both ends.
 *
 * The padding is what lets a multi-word term be found with a plain substring
 * search: « impot » can then never match inside « impotent ».
 */
export function stems(text) {
  const found = normalise(text).match(WORD) ?? [];
  return ` ${found.map(lemmatise).join(' ')} `;
}

/**
 * Every topic whose terms appear in the article, best supported first.
 *
 * `vocabulary` maps a topic name to the terms that stand for it, single or
 * multi-word. Keeping it a parameter is the point of this rung: the taxonomy
 * belongs to whoever edits the articles, not to the code.
 *
 * `minTerms` is how many distinct terms a topic needs before it is claimed.
 * Raise it when a single passing mention is not enough to file an article.
 */
export function tag(article, vocabulary, minTerms = 1) {
  const haystack = stems(article);
  const hits = new Map();
  for (const [topic, terms] of Object.entries(vocabulary)) {
    const found = terms.filter((term) => haystack.includes(stems(term))).length;
    if (found >= minTerms) hits.set(topic, found);
  }
  // Best supported first, ties in alphabetical order: a tagging run has to be
  // replayable, and a set has no order to replay.
  return [...hits.keys()].sort((a, b) => hits.get(b) - hits.get(a) || a.localeCompare(b));
}

"""
Tag articles with a one-versus-rest classifier over TF-IDF features.

Rung N1. The controlled vocabulary of N0 sees the words an editor listed. This
sees the words that go with a topic in the articles you have already tagged —
« bureau », « visioconférence » and « domicile » end up carrying the remote
work topic, although no editor would ever have written them in a term list.

One classifier per topic, each answering its own yes-or-no question. That is
what one-versus-rest means, and it is what keeps the tagging multi-label: the
topics do not compete for a single winner, so an article can come back with
three tags, or with none.

The cost of this rung is not the code, which is below. It is the labelled
corpus: a few hundred articles someone has to tag by hand, and tag again every
time the taxonomy moves.
"""

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import MultiLabelBinarizer


def train(articles: list[str], topics_per_article: list[list[str]]) -> dict:
    """
    `topics_per_article[i]` is the list of topics of `articles[i]`, possibly
    empty. An untagged article is a useful negative example, not a gap.
    """
    binariser = MultiLabelBinarizer()
    matrix = binariser.fit_transform(topics_per_article)
    pipeline = make_pipeline(
        # Word unigrams and bigrams: « à distance » says something that
        # « distance » alone does not.
        TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True),
        OneVsRestClassifier(LogisticRegression(C=4.0, max_iter=1000)),
    )
    pipeline.fit(articles, matrix)
    return {"pipeline": pipeline, "topics": list(binariser.classes_)}


def score(model: dict, article: str) -> dict[str, float]:
    """One probability per topic, each independent of the others."""
    probabilities = model["pipeline"].predict_proba([article])[0]
    return {topic: float(p) for topic, p in zip(model["topics"], probabilities)}


def tag(model: dict, article: str, threshold: float = 0.5) -> list[str]:
    """
    The topics above the threshold, best first.

    The threshold is yours to set, and it is the only dial here. Move it
    towards 1 when a wrong tag is worse than a missing one, towards 0 when an
    editor reviews the list anyway and would rather see one topic too many.
    """
    scored = score(model, article)
    kept = [topic for topic, value in scored.items() if value >= threshold]
    return sorted(kept, key=lambda topic: (-scored[topic], topic))

/**
 * Tag articles with a one-versus-rest classifier over TF-IDF features.
 *
 * Rung N1. The controlled vocabulary of N0 sees the words an editor listed.
 * This sees the words that go with a topic in the articles you have already
 * tagged — « bureau », « visioconférence » and « domicile » end up carrying
 * the remote work topic, although no editor would ever have written them in a
 * term list.
 *
 * One classifier per topic, each answering its own yes-or-no question. That is
 * what one-versus-rest means, and it is what keeps the tagging multi-label:
 * the topics do not compete for a single winner, so an article can come back
 * with three tags, or with none.
 *
 * Written out in full rather than pulled from a library, because TF-IDF and
 * logistic regression fit on a page between them. The cost of this rung is not
 * the code below: it is the labelled corpus someone has to tag by hand, and
 * tag again every time the taxonomy moves.
 */

const TOKEN = /[\p{L}\p{N}_]{2,}/gu;

/** Unigrams and bigrams: « à distance » says more than « distance » alone. */
function terms(text) {
  const words = text.toLowerCase().match(TOKEN) ?? [];
  const pairs = words.slice(1).map((word, i) => `${words[i]} ${word}`);
  return [...words, ...pairs];
}

/** Sublinear term frequency times inverse document frequency, L2 normalised. */
function vectorise(termList, vocabulary, idf) {
  const counts = new Map();
  for (const term of termList) counts.set(term, (counts.get(term) ?? 0) + 1);
  const vector = new Float64Array(idf.length);
  for (const [term, count] of counts) {
    const column = vocabulary.get(term);
    if (column !== undefined) vector[column] = (1 + Math.log(count)) * idf[column];
  }
  const norm = Math.hypot(...vector);
  if (norm) for (let j = 0; j < vector.length; j += 1) vector[j] /= norm;
  return vector;
}

/** One binary classifier, trained by gradient descent one article at a time. */
function fit(rows, labels, { epochs, rate, decay }) {
  const weights = new Float64Array(rows[0].length);
  let bias = 0;
  for (let epoch = 0; epoch < epochs; epoch += 1) {
    for (let i = 0; i < rows.length; i += 1) {
      let z = bias;
      for (let j = 0; j < weights.length; j += 1) z += weights[j] * rows[i][j];
      const error = 1 / (1 + Math.exp(-z)) - labels[i];
      // The decay term is the L2 penalty: without it a corpus this small is
      // memorised, and every article comes back scored one or zero.
      for (let j = 0; j < weights.length; j += 1) {
        weights[j] -= rate * (error * rows[i][j] + decay * weights[j]);
      }
      bias -= rate * error;
    }
  }
  return { weights, bias };
}

/**
 * `topicsPerArticle[i]` is the list of topics of `articles[i]`, possibly
 * empty. An untagged article is a useful negative example, not a gap.
 */
export function train(articles, topicsPerArticle, { epochs = 600, rate = 0.5, decay = 0.005 } = {}) {
  const documents = articles.map(terms);
  const frequencies = new Map();
  for (const document of documents) {
    for (const term of new Set(document)) frequencies.set(term, (frequencies.get(term) ?? 0) + 1);
  }
  const vocabulary = new Map([...frequencies.keys()].map((term, column) => [term, column]));
  const idf = [...frequencies.values()].map((df) => Math.log((1 + documents.length) / (1 + df)) + 1);
  const rows = documents.map((document) => vectorise(document, vocabulary, idf));
  const topics = [...new Set(topicsPerArticle.flat())].sort();
  const classifiers = topics.map((topic) => {
    const labels = topicsPerArticle.map((list) => (list.includes(topic) ? 1 : 0));
    return fit(rows, labels, { epochs, rate, decay });
  });
  return { vocabulary, idf, topics, classifiers };
}

/** One probability per topic, each independent of the others. */
export function score(model, article) {
  const vector = vectorise(terms(article), model.vocabulary, model.idf);
  const scored = {};
  model.topics.forEach((topic, index) => {
    const { weights, bias } = model.classifiers[index];
    let z = bias;
    for (let j = 0; j < weights.length; j += 1) z += weights[j] * vector[j];
    scored[topic] = 1 / (1 + Math.exp(-z));
  });
  return scored;
}

/**
 * The topics above the threshold, best first.
 *
 * The threshold is yours to set, and it is the only dial here. Move it towards
 * 1 when a wrong tag is worse than a missing one, towards 0 when an editor
 * reviews the list anyway and would rather see one topic too many.
 */
export function tag(model, article, threshold = 0.5) {
  const scored = score(model, article);
  return model.topics
    .filter((topic) => scored[topic] >= threshold)
    .sort((a, b) => scored[b] - scored[a] || a.localeCompare(b));
}

"""
Zero-shot tagging: compare the article with the labels themselves.

Rung N2. N1 needed a labelled corpus, and a new topic meant labelling it all
again. Here a topic is only a short description, encoded like any other text:
adding one costs a line, and the first article can be tagged the same day.

That is the whole appeal, and it is real. What it costs is a model file to
ship and keep in sync, a warm process to hold it, and a score that is a cosine
rather than a probability — see the test for what that means when you have to
pick one threshold for every topic.
"""

from __future__ import annotations

MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"


def load_encoder(name: str = MODEL_NAME):
    """The real encoder: fetched once, then held in memory and run locally."""
    from sentence_transformers import SentenceTransformer  # pragma: no cover

    return SentenceTransformer(name)


def build_labeller(topics: dict[str, str], encoder=None) -> dict:
    """
    Encode the topic descriptions once, and keep the encoder for the articles.

    `topics` maps a topic name to the sentence that stands for it. Write it as
    a human would say it out loud: the encoder was trained on sentences, and a
    bare keyword gives it very little to work with.

    `encoder` is injected so this can be tested without loading a model. Left
    alone, it is the real one above.
    """
    encoder = load_encoder() if encoder is None else encoder
    names = list(topics)
    vectors = [_unit(v) for v in encoder.encode([topics[name] for name in names])]
    return {"topics": names, "vectors": vectors, "encoder": encoder}


def score(labeller: dict, article: str) -> dict[str, float]:
    """The cosine between the article and each topic description."""
    vector = _unit(labeller["encoder"].encode([article])[0])
    return {
        topic: _dot(known, vector)
        for topic, known in zip(labeller["topics"], labeller["vectors"])
    }


def tag(labeller: dict, article: str, threshold: float = 0.3) -> list[str]:
    """
    Every topic whose description is close enough, best first.

    Multi-label falls out of the shape of the thing: each topic is compared
    with the article on its own, so several can pass, or none.

    The threshold is not a probability. It is a cosine, it has no calibrated
    meaning, and the only way to set it is to try it on articles you have
    already tagged by hand — which is a labelled corpus, the very thing this
    rung was supposed to save you.
    """
    scored = score(labeller, article)
    kept = [topic for topic, value in scored.items() if value >= threshold]
    return sorted(kept, key=lambda topic: (-scored[topic], topic))


def _unit(vector) -> list[float]:
    """Cosine similarity is a dot product once both sides have length one."""
    values = [float(v) for v in vector]
    norm = sum(v * v for v in values) ** 0.5
    return [v / norm for v in values] if norm else values


def _dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))

/**
 * Zero-shot tagging: compare the article with the labels themselves.
 *
 * Rung N2. N1 needed a labelled corpus, and a new topic meant labelling it all
 * again. Here a topic is only a short description, encoded like any other
 * text: adding one costs a line, and the first article can be tagged the same
 * day.
 *
 * That is the whole appeal, and it is real. What it costs is a model file to
 * ship and keep in sync, a warm process to hold it, and a score that is a
 * cosine rather than a probability — see the test for what that means when you
 * have to pick one threshold for every topic.
 */

export const MODEL_NAME = 'Xenova/paraphrase-multilingual-MiniLM-L12-v2';

/** The real encoder: fetched once, then held in memory and run locally. */
export async function loadEncoder(name = MODEL_NAME) {
  const { pipeline } = await import('@xenova/transformers');
  const extract = await pipeline('feature-extraction', name);
  return { encode: async (texts) => (await extract(texts, { pooling: 'mean' })).tolist() };
}

/**
 * Encode the topic descriptions once, and keep the encoder for the articles.
 *
 * `topics` maps a topic name to the sentence that stands for it. Write it as a
 * human would say it out loud: the encoder was trained on sentences, and a
 * bare keyword gives it very little to work with.
 *
 * `encoder` is injected so this can be tested without loading a model. Left
 * alone, it is the real one above.
 */
export async function buildLabeller(topics, encoder) {
  const model = encoder ?? (await loadEncoder());
  const names = Object.keys(topics);
  const vectors = (await model.encode(names.map((name) => topics[name]))).map(unit);
  return { topics: names, vectors, encoder: model };
}

/** The cosine between the article and each topic description. */
export async function score(labeller, article) {
  const vector = unit((await labeller.encoder.encode([article]))[0]);
  const scored = {};
  labeller.topics.forEach((topic, index) => {
    scored[topic] = dot(labeller.vectors[index], vector);
  });
  return scored;
}

/**
 * Every topic whose description is close enough, best first.
 *
 * Multi-label falls out of the shape of the thing: each topic is compared with
 * the article on its own, so several can pass, or none.
 *
 * The threshold is not a probability. It is a cosine, it has no calibrated
 * meaning, and the only way to set it is to try it on articles you have
 * already tagged by hand — which is a labelled corpus, the very thing this
 * rung was supposed to save you.
 */
export async function tag(labeller, article, threshold = 0.3) {
  const scored = await score(labeller, article);
  return labeller.topics
    .filter((topic) => scored[topic] >= threshold)
    .sort((a, b) => scored[b] - scored[a] || a.localeCompare(b));
}

/** Cosine similarity is a dot product once both sides have length one. */
function unit(vector) {
  let sum = 0;
  for (const v of vector) sum += v * v;
  const norm = Math.sqrt(sum);
  return norm ? vector.map((v) => v / norm) : [...vector];
}

function dot(a, b) {
  let total = 0;
  for (let i = 0; i < a.length; i += 1) total += a[i] * b[i];
  return total;
}

"""
Tag articles by asking a general-purpose model.

Rung N3. This is the option people reach for first, and on this entry it is
the only one that needs neither a term list, nor a labelled corpus, nor a
model file. 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 size, retry on
failure, parse an answer that is only probably valid JSON, and keep only the
topics that exist in your taxonomy. That last one is not optional — a model
will happily invent a plausible topic — and it is why the vocabulary stays a
parameter here, exactly as it was on the bottom rung.

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.
"""

from __future__ import annotations

import json

PROMPT = (
    "Tag the article below with the topics it covers.\n"
    "Choose only from this list, and answer with the spellings given:\n"
    "{topics}\n"
    "An article may cover several topics, or none at all.\n"
    "Answer with JSON only: a list of topic names, empty if none apply.\n\n"
    "Article:\n{article}"
)

MAX_CHARACTERS = 12000


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


def tag(article: str, topics: list[str], client=None, *, attempts: int = 3) -> list[str]:
    """
    The topics of the article, in the order of the taxonomy.

    `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(article) > MAX_CHARACTERS:
        raise ValueError(f"article longer than {MAX_CHARACTERS} characters")

    reported = _ask(client, article, topics, attempts)

    # Keep only what the taxonomy knows, and answer in the taxonomy's own
    # order. A model that invents « actualité juridique » must not create a
    # topic in your database, and two identical calls must file an article the
    # same way twice.
    answered = {name.strip().lower() for name in reported if isinstance(name, str)}
    return [topic for topic in topics if topic.lower() in answered]


def _ask(client, article: str, topics: list[str], attempts: int) -> list:
    prompt = PROMPT.format(
        topics="\n".join(f"- {topic}" for topic in topics), article=article
    )
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            answer = client.complete(prompt=prompt, temperature=0)
            parsed = json.loads(answer)
            if isinstance(parsed, list):
                return parsed
            last_error = ValueError("the model answered something that is not a list")
        except Exception as error:  # noqa: BLE001 - any provider failure is retried
            last_error = error
    raise TaggingUnavailable(str(last_error))

/**
 * Tag articles by asking a general-purpose model.
 *
 * Rung N3. This is the option people reach for first, and on this entry it is
 * the only one that needs neither a term list, nor a labelled corpus, nor a
 * model file. 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 size, retry on
 * failure, parse an answer that is only probably valid JSON, and keep only the
 * topics that exist in your taxonomy. That last one is not optional — a model
 * will happily invent a plausible topic — and it is why the vocabulary stays a
 * parameter here, exactly as it was on the bottom rung.
 *
 * 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.
 */

const PROMPT = [
  'Tag the article below with the topics it covers.',
  'Choose only from this list, and answer with the spellings given:',
  '{topics}',
  'An article may cover several topics, or none at all.',
  'Answer with JSON only: a list of topic names, empty if none apply.',
  '',
  'Article:',
].join('\n');

export const MAX_CHARACTERS = 12000;

export class TaggingUnavailable extends Error {}

/**
 * The topics of the article, in the order of the taxonomy.
 *
 * @param {string} article
 * @param {string[]} topics the controlled vocabulary, as on rung N0
 * @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 tag(article, topics, { 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 (article.length > MAX_CHARACTERS) {
    throw new RangeError(`article longer than ${MAX_CHARACTERS} characters`);
  }

  const reported = await ask(client, article, topics, attempts);

  // Keep only what the taxonomy knows, and answer in the taxonomy's own order.
  // A model that invents « actualité juridique » must not create a topic in
  // your database, and two identical calls must file an article the same way
  // twice.
  const answered = new Set(
    reported.filter((name) => typeof name === 'string').map((name) => name.trim().toLowerCase()),
  );
  return topics.filter((topic) => answered.has(topic.toLowerCase()));
}

async function ask(client, article, topics, attempts) {
  const prompt = `${PROMPT.replace('{topics}', topics.map((t) => `- ${t}`).join('\n'))}\n${article}`;
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      const answer = await client.complete({
        prompt,
        // Temperature zero, because a taxonomy that changes between two
        // identical calls is not a taxonomy.
        temperature: 0,
      });
      const parsed = JSON.parse(answer);
      if (Array.isArray(parsed)) return parsed;
      lastError = new Error('the model answered something that is not a list');
    } catch (error) {
      lastError = error;
    }
  }
  throw new TaggingUnavailable(String(lastError));
}