Summarise a long document

"""
Summarise a long document by choosing its own best sentences.

Rung N0. Extractive: every sentence of the summary appears verbatim in the
source, because this code never writes a word, it only selects.

Two classical signals, and nothing else.

Term frequency. A word the document keeps coming back to is what the document
is about, so a sentence dense in such words carries more of the subject than a
sentence made of connectives. Dividing by the length of the sentence measures
density rather than volume, which stops a long sentence from winning by size.

Position. An author states the subject early. A small bonus that decays with
the rank of the sentence encodes that habit, without handing the summary to
the opening paragraph outright.
"""

import re
import unicodedata

# A sentence ends at a full stop, question or exclamation mark followed by
# whitespace, or at a line break, so a transcript or a list without final
# punctuation is still cut into lines. Abbreviations will fool this, and so
# will prose hard-wrapped at a fixed width, which has to be unwrapped first; a
# real corpus needs a better splitter, a separate problem from choosing.
SENTENCE_END = re.compile(r"(?<=[.!?])\s+|\s*\n\s*")

# Letters and digits only: `[^\W_]` is `\w` without the underscore, so
# accented words survive and punctuation does not.
WORD = re.compile(r"[^\W_]+")


def words(sentence: str) -> list[str]:
    """Lowercase words, composed first: a decomposed `é` would split the word."""
    return WORD.findall(unicodedata.normalize("NFC", sentence).lower())

# Words too common to say anything about the subject of a document.
STOPWORDS = frozenset(
    "a an and are as at be been but by for from had has have in into is it its "
    "of on or that the their there they this to was were which will with".split()
)

# How much the opening of the document is worth. Large enough to break a tie
# between two equally dense sentences, too small to win on its own.
LEAD_BONUS = 0.15


def split_sentences(text: str) -> list[str]:
    """Cut the document into sentences, dropping empty ones."""
    parts = SENTENCE_END.split(text.strip())
    return [part.strip() for part in parts if part.strip()]


def _term_weights(sentences: list[str]) -> dict[str, float]:
    """Count content words, then scale so the most frequent one weighs one."""
    counts: dict[str, int] = {}
    for sentence in sentences:
        for word in words(sentence):
            if len(word) > 2 and word not in STOPWORDS:
                counts[word] = counts.get(word, 0) + 1
    if not counts:
        return {}
    most = max(counts.values())
    return {word: count / most for word, count in counts.items()}


def score_sentences(sentences: list[str]) -> list[float]:
    """Density in the document's own vocabulary, plus the position bonus."""
    weights = _term_weights(sentences)
    scores = []
    for index, sentence in enumerate(sentences):
        found = words(sentence)
        # Accumulated in a plain loop rather than with `sum`, which since
        # Python 3.12 compensates rounding error on floats. That is the better
        # answer, but it is not the answer JavaScript gives, and the two
        # versions of this snippet have to rank sentences identically.
        total = 0.0
        for word in found:
            total += weights.get(word, 0.0)
        density = total / len(found) if found else 0.0
        scores.append(density + LEAD_BONUS / (index + 1))
    return scores


def summarise(text: str, max_sentences: int = 3) -> str:
    """
    Return the best sentences, in the order the document puts them.

    Ordering the summary by score would read as a list of quotations. Keeping
    document order keeps the sequence the author chose, which is the only part
    of the argument an extractive summary can preserve.
    """
    if max_sentences < 0:
        raise ValueError("max_sentences cannot be negative")
    sentences = split_sentences(text)
    scores = score_sentences(sentences)
    # Sorting is stable, so two identical scores keep their document order.
    ranked = sorted(range(len(sentences)), key=lambda i: -scores[i])
    chosen = sorted(ranked[:max_sentences])
    return " ".join(sentences[index] for index in chosen)

/**
 * Summarise a long document by choosing its own best sentences.
 *
 * Rung N0. Extractive: every sentence of the summary appears verbatim in the
 * source, because this code never writes a word, it only selects.
 *
 * Two classical signals, and nothing else.
 *
 * Term frequency. A word the document keeps coming back to is what the
 * document is about, so a sentence dense in such words carries more of the
 * subject than a sentence made of connectives. Dividing by the length of the
 * sentence measures density rather than volume, which stops a long sentence
 * from winning by size.
 *
 * Position. An author states the subject early. A small bonus that decays
 * with the rank of the sentence encodes that habit, without handing the
 * summary to the opening paragraph outright.
 */

// A sentence ends at a full stop, question or exclamation mark followed by
// whitespace, or at a line break, so a transcript or a list without final
// punctuation is still cut into lines. Abbreviations will fool this, and so
// will prose hard-wrapped at a fixed width, which has to be unwrapped first; a
// real corpus needs a better splitter, a separate problem from choosing.
const SENTENCE_END = /(?<=[.!?])\s+|\s*\n\s*/;

// Letters and digits only, so accented words survive and punctuation does
// not. The Python counterpart writes the same class as `[^\W_]`.
const WORD = /[\p{L}\p{N}]+/gu;

/** Lowercase words, composed first: a decomposed `é` would split the word. */
function words(sentence) {
  return sentence.normalize('NFC').toLowerCase().match(WORD) ?? [];
}

// Words too common to say anything about the subject of a document.
const STOPWORDS = new Set(
  ('a an and are as at be been but by for from had has have in into is it its ' +
    'of on or that the their there they this to was were which will with').split(' '),
);

// How much the opening of the document is worth. Large enough to break a tie
// between two equally dense sentences, too small to win on its own.
const LEAD_BONUS = 0.15;

/** Cut the document into sentences, dropping empty ones. */
export function splitSentences(text) {
  return text
    .trim()
    .split(SENTENCE_END)
    .map((part) => part.trim())
    .filter(Boolean);
}

/** Count content words, then scale so the most frequent one weighs one. */
function termWeights(sentences) {
  const counts = new Map();
  for (const sentence of sentences) {
    for (const word of words(sentence)) {
      if (word.length > 2 && !STOPWORDS.has(word)) {
        counts.set(word, (counts.get(word) ?? 0) + 1);
      }
    }
  }
  if (counts.size === 0) return new Map();
  const most = Math.max(...counts.values());
  return new Map([...counts].map(([word, count]) => [word, count / most]));
}

/** Density in the document's own vocabulary, plus the position bonus. */
export function scoreSentences(sentences) {
  const weights = termWeights(sentences);
  return sentences.map((sentence, index) => {
    const found = words(sentence);
    let total = 0;
    for (const word of found) total += weights.get(word) ?? 0;
    const density = found.length ? total / found.length : 0;
    return density + LEAD_BONUS / (index + 1);
  });
}

/**
 * Return the best sentences, in the order the document puts them.
 *
 * Ordering the summary by score would read as a list of quotations. Keeping
 * document order keeps the sequence the author chose, which is the only part
 * of the argument an extractive summary can preserve.
 */
export function summarise(text, maxSentences = 3) {
  if (maxSentences < 0) throw new RangeError('maxSentences cannot be negative');
  const sentences = splitSentences(text);
  const scores = scoreSentences(sentences);
  // Sorting is stable, so two identical scores keep their document order.
  const ranked = sentences.map((_, index) => index).sort((a, b) => scores[b] - scores[a]);
  const chosen = ranked.slice(0, maxSentences).sort((a, b) => a - b);
  return chosen.map((index) => sentences[index]).join(' ');
}

"""
Score sentences with a classifier trained on surface features.

Rung N1. Same extractive shape as N0 — the summary is still made of the
document's own sentences — but the weights are learnt instead of guessed.

N0 fixes the trade between position and term density by hand, once, for every
document in the world. That constant is a guess. Here, documents whose summary
sentences someone has ticked off set the weights of five features instead, and
the model learns nothing those five cannot express. Position is coded as
1 / (rank + 1): the model can learn that early sentences count for more or for
less, never that the last one counts. A closing sentence gains only through
what it carries, a figure or a cue word such as `overall`.

The features are deliberately surface-level. They describe where a sentence
sits and what it looks like, not what it means. That is the ceiling of this
rung, and the reason the entry does not stop here.
"""

import re
import unicodedata

import numpy as np
from sklearn.linear_model import LogisticRegression

# A full stop, question or exclamation mark followed by whitespace, or a line
# break: the same splitter as N0, with the same limits.
SENTENCE_END = re.compile(r"(?<=[.!?])\s+|\s*\n\s*")
WORD = re.compile(r"[^\W_]+")

STOPWORDS = frozenset(
    "a an and are as at be been but by for from had has have in into is it its "
    "of on or that the their there they this to was were which will with".split()
)

# Words an author uses when about to state the point of what came before.
CUES = frozenset(
    "overall therefore total conclusion result summary altogether finally "
    "consequently".split()
)

# Sentences longer than this are already long; the feature saturates rather
# than letting one outlier stretch the scale for every other sentence.
LONG_SENTENCE = 25


def words(sentence: str) -> list[str]:
    """Lowercase words, composed first: a decomposed `é` would split the word."""
    return WORD.findall(unicodedata.normalize("NFC", sentence).lower())


def split_sentences(text: str) -> list[str]:
    """Cut the document into sentences, dropping empty ones."""
    parts = SENTENCE_END.split(text.strip())
    return [part.strip() for part in parts if part.strip()]


def sentence_features(sentences: list[str], index: int) -> list[float]:
    """
    Five things a reader notices before reading: where the sentence sits, how
    long it is, whether it carries a figure, whether it announces a conclusion,
    and how much of the opening it repeats.
    """
    sentence = sentences[index]
    found = words(sentence)
    opening = set(words(sentences[0]))
    content = [w for w in found if len(w) > 2 and w not in STOPWORDS]
    shared = sum(1 for word in content if word in opening)
    return [
        1.0 / (index + 1),
        min(len(found) / LONG_SENTENCE, 1.0),
        1.0 if any(c in "0123456789" for c in sentence) else 0.0,
        1.0 if any(word in CUES for word in found) else 0.0,
        shared / len(content) if content else 0.0,
    ]


def train(documents: list[list[str]], labels: list[list[int]]) -> LogisticRegression:
    """
    `documents` are lists of sentences; `labels[d][i]` is 1 when sentence `i`
    of document `d` belongs in the summary.

    The ticked sentences are the entire training set. Nothing here needs a GPU.
    """
    rows, targets = [], []
    for sentences, marks in zip(documents, labels):
        for index in range(len(sentences)):
            rows.append(sentence_features(sentences, index))
            targets.append(marks[index])
    return LogisticRegression(max_iter=2000).fit(np.array(rows), np.array(targets))


def summarise(model: LogisticRegression, text: str, max_sentences: int = 3) -> str:
    """Return the best-scored sentences, in the order the document puts them."""
    if max_sentences < 0:
        raise ValueError("max_sentences cannot be negative")
    sentences = split_sentences(text)
    if not sentences:
        return ""
    rows = np.array([sentence_features(sentences, i) for i in range(len(sentences))])
    scores = model.predict_proba(rows)[:, 1]
    # Sorting is stable, so two identical scores keep their document order.
    ranked = sorted(range(len(sentences)), key=lambda i: -scores[i])
    chosen = sorted(ranked[:max_sentences])
    return " ".join(sentences[index] for index in chosen)

/**
 * Score sentences with a classifier trained on surface features.
 *
 * Rung N1. Same extractive shape as N0 — the summary is still made of the
 * document's own sentences — but the weights are learnt instead of guessed.
 *
 * N0 fixes the trade between position and term density by hand, once, for
 * every document in the world. That constant is a guess. Here, documents whose
 * summary sentences someone has ticked off set the weights of five features
 * instead, and the model learns nothing those five cannot express. Position is
 * coded as 1 / (rank + 1): the model can learn that early sentences count for
 * more or for less, never that the last one counts. A closing sentence gains
 * only through what it carries, a figure or a cue word such as `overall`.
 *
 * The features are deliberately surface-level. They describe where a sentence
 * sits and what it looks like, not what it means. That is the ceiling of this
 * rung, and the reason the entry does not stop here.
 *
 * Logistic regression is written out rather than pulled from a package, so
 * the snippet needs nothing but Node. The Python version of this
 * snippet calls scikit-learn; the objective minimised below is the same one,
 * so both rank a document's sentences alike.
 */

// A full stop, question or exclamation mark followed by whitespace, or a line
// break: the same splitter as N0, with the same limits.
const SENTENCE_END = /(?<=[.!?])\s+|\s*\n\s*/;
const WORD = /[\p{L}\p{N}]+/gu;

/** Lowercase words, composed first: a decomposed `é` would split the word. */
function words(sentence) {
  return sentence.normalize('NFC').toLowerCase().match(WORD) ?? [];
}

const STOPWORDS = new Set(
  ('a an and are as at be been but by for from had has have in into is it its ' +
    'of on or that the their there they this to was were which will with').split(' '),
);

// Words an author uses when about to state the point of what came before.
const CUES = new Set(
  ('overall therefore total conclusion result summary altogether finally ' +
    'consequently').split(' '),
);

// Sentences longer than this are already long; the feature saturates rather
// than letting one outlier stretch the scale for every other sentence.
const LONG_SENTENCE = 25;

/** Cut the document into sentences, dropping empty ones. */
export function splitSentences(text) {
  return text
    .trim()
    .split(SENTENCE_END)
    .map((part) => part.trim())
    .filter(Boolean);
}

/**
 * Five things a reader notices before reading: where the sentence sits, how
 * long it is, whether it carries a figure, whether it announces a conclusion,
 * and how much of the opening it repeats.
 */
export function sentenceFeatures(sentences, index) {
  const sentence = sentences[index];
  const found = words(sentence);
  const opening = new Set(words(sentences[0]));
  const content = found.filter((w) => w.length > 2 && !STOPWORDS.has(w));
  const shared = content.filter((word) => opening.has(word)).length;
  return [
    1 / (index + 1),
    Math.min(found.length / LONG_SENTENCE, 1),
    /[0-9]/.test(sentence) ? 1 : 0,
    found.some((word) => CUES.has(word)) ? 1 : 0,
    content.length ? shared / content.length : 0,
  ];
}

/**
 * `documents` are lists of sentences; `labels[d][i]` is 1 when sentence `i` of
 * document `d` belongs in the summary.
 *
 * Full-batch gradient descent on the mean log-loss, plus the same L2 penalty
 * scikit-learn applies by default. The bias is left unpenalised, as it is
 * there too. The objective is strictly convex, so enough steps land on the one
 * optimum whichever language walks towards it.
 */
export function train(documents, labels, { epochs = 4000, rate = 1, strength = 1 } = {}) {
  const rows = [];
  const targets = [];
  for (const [d, sentences] of documents.entries()) {
    for (let i = 0; i < sentences.length; i += 1) {
      rows.push(sentenceFeatures(sentences, i));
      targets.push(labels[d][i]);
    }
  }
  const width = rows[0].length;
  const weights = new Array(width).fill(0);
  let bias = 0;
  const penalty = 1 / (strength * rows.length);

  for (let epoch = 0; epoch < epochs; epoch += 1) {
    const gradient = new Array(width).fill(0);
    let biasGradient = 0;
    for (let i = 0; i < rows.length; i += 1) {
      let z = bias;
      for (let j = 0; j < width; j += 1) z += weights[j] * rows[i][j];
      const error = 1 / (1 + Math.exp(-z)) - targets[i];
      for (let j = 0; j < width; j += 1) gradient[j] += (error * rows[i][j]) / rows.length;
      biasGradient += error / rows.length;
    }
    for (let j = 0; j < width; j += 1) weights[j] -= rate * (gradient[j] + penalty * weights[j]);
    bias -= rate * biasGradient;
  }
  return { weights, bias };
}

/** Probability that a sentence belongs in the summary. */
function score(model, features) {
  let z = model.bias;
  for (let j = 0; j < features.length; j += 1) z += model.weights[j] * features[j];
  return 1 / (1 + Math.exp(-z));
}

/** Return the best-scored sentences, in the order the document puts them. */
export function summarise(model, text, maxSentences = 3) {
  if (maxSentences < 0) throw new RangeError('maxSentences cannot be negative');
  const sentences = splitSentences(text);
  if (sentences.length === 0) return '';
  const scores = sentences.map((_, i) => score(model, sentenceFeatures(sentences, i)));
  // Sorting is stable, so two identical scores keep their document order.
  const ranked = sentences.map((_, index) => index).sort((a, b) => scores[b] - scores[a]);
  const chosen = ranked.slice(0, maxSentences).sort((a, b) => a - b);
  return chosen.map((index) => sentences[index]).join(' ');
}

"""
Summarise a long document with a self-hosted abstractive model.

Rung N2. The first rung that writes. N0 and N1 choose sentences and can only
ever return what the author already wrote; this one produces sentences that
were not in the document.

That gain has a price, and the price is in this file rather than in the model.
A sequence-to-sequence summariser reads a fixed-size window. A long document
does not fit, so it has to be cut at sentence boundaries, summarised piece by
piece, and the notes summarised again, in as many passes as it takes to fit
one window. Every one of those calls can fail or come back empty, and a
summary that is silently half a document is worse than no summary at all.

What this file cannot do, at any price: check that what the model wrote is
what the document said. Nothing in this plumbing can. See the test.

Nor does it check the language. The checkpoint named below declares
`language: en` on its card and was fine-tuned on English newswire; on a French
note it will write something, and nothing here will stop it. A multilingual
summariser exists — mT5 fine-tuned on XL-Sum, forty-five languages, French
among them — under CC BY-SA-NC 4.0, which rules out commercial use. Read the
card before you pick, and read it for the language first.
"""

from __future__ import annotations

import functools
import re

# A full stop, question or exclamation mark followed by whitespace, or a line
# break: a transcript without final punctuation is still cut into lines.
SENTENCE_END = re.compile(r"(?<=[.!?])\s+|\s*\n\s*")

# The same weights in both languages: JavaScript loads their ONNX conversion.
# English only: its card says `language: en`, and it is distilled from a model
# fine-tuned on CNN/DailyMail and XSum, which are English newswire.
MODEL_NAME = "sshleifer/distilbart-cnn-12-6"

# What one call to the model is given to read. The window is 1,024 tokens, and
# characters are not tokens: with this model's tokenizer, 3,000 characters of
# French prose come to 1,091 tokens, 2,000 to 727. A piece that still
# overflows, code or a table, is refused by the model below, not truncated.
CHUNK_CHARACTERS = 2000

# What the whole function is allowed to read. A self-hosted model costs
# machine time rather than money, but a document nobody meant to send is still
# better refused than churned through in silence.
MAX_CHARACTERS = 200_000


class SummaryUnavailable(Exception):
    """The model failed, or returned nothing usable."""


class LocalSummariser:
    """The real model, loaded from local weights and run on this machine."""

    def __init__(self, name: str = MODEL_NAME) -> None:  # pragma: no cover - loads weights
        # transformers 5 removed the "summarization" pipeline: call the model itself.
        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

        self._tokenizer = AutoTokenizer.from_pretrained(name)
        self._model = AutoModelForSeq2SeqLM.from_pretrained(name)
        # The generation settings the checkpoint was published with.
        self._settings = self._model.config.task_specific_params["summarization"]

    def generate(self, text: str) -> str:  # pragma: no cover - loads weights
        inputs = self._tokenizer(text, return_tensors="pt")
        # Refused rather than truncated: a summary of the start of a piece
        # reads exactly like a summary of all of it.
        if inputs["input_ids"].shape[1] > self._tokenizer.model_max_length:
            raise ValueError("this piece is longer than the model's window")
        output = self._model.generate(**inputs, **self._settings)
        return self._tokenizer.decode(output[0], skip_special_tokens=True)


@functools.cache
def default_model() -> LocalSummariser:  # pragma: no cover - loads weights
    """Loaded once and kept for the life of the process: the loading is the slow part."""
    return LocalSummariser()


def pack(units: list[str], size: int = CHUNK_CHARACTERS) -> list[str]:
    """Join consecutive units into pieces of at most `size` characters."""
    pieces: list[str] = []
    current = ""
    for unit in units:
        if current and len(current) + 1 + len(unit) > size:
            pieces.append(current)
            current = unit
        else:
            current = f"{current} {unit}".strip()
    if current:
        pieces.append(current)
    return pieces


def chunk(text: str, size: int = CHUNK_CHARACTERS) -> list[str]:
    """
    Cut the document into pieces that fit the model's window, at sentence
    boundaries.

    A sentence longer than the window on its own is passed whole. Cutting
    mid-sentence would hand the model half a clause, which is a worse thing to
    summarise.
    """
    sentences = (part.strip() for part in SENTENCE_END.split(text.strip()))
    return pack([sentence for sentence in sentences if sentence], size)


def summarise(text: str, model=None, *, attempts: int = 2) -> str:
    """
    `model` is injected so this function can be tested without loading weights.
    In production it defaults to the real model.
    """
    if len(text) > MAX_CHARACTERS:
        raise ValueError(f"document longer than {MAX_CHARACTERS} characters")
    pieces = chunk(text)
    if not pieces:
        return ""
    model = model or default_model()

    while True:
        notes = [_generate(model, piece, attempts) for piece in pieces]
        if len(notes) == 1:
            return notes[0]
        # Next pass: the model reads back its own notes, packed to the window
        # like the document was. It never reads the document again.
        packed = pack(notes)
        if len(packed) >= len(pieces):
            raise SummaryUnavailable("the notes are no shorter than what they summarise")
        pieces = packed


def _generate(model, text: str, attempts: int) -> str:
    last_error: Exception | None = None
    for _ in range(attempts):
        try:
            written = (model.generate(text) or "").strip()
            if written:
                return written
            # An empty answer is a failure, not a summary. Returning it would
            # leave a hole in the middle of the document with no trace.
            last_error = ValueError("the model returned an empty summary")
        except Exception as error:  # noqa: BLE001 - any model failure is retried
            last_error = error
    raise SummaryUnavailable(str(last_error))

/**
 * Summarise a long document with a self-hosted abstractive model.
 *
 * Rung N2. The first rung that writes. N0 and N1 choose sentences and can only
 * ever return what the author already wrote; this one produces sentences that
 * were not in the document.
 *
 * That gain has a price, and the price is in this file rather than in the
 * model. A sequence-to-sequence summariser reads a fixed-size window. A long
 * document does not fit, so it has to be cut at sentence boundaries,
 * summarised piece by piece, and the notes summarised again, in as many passes
 * as it takes to fit one window. Every one of those calls can fail or come
 * back empty, and a summary that is silently half a document is worse than no
 * summary at all.
 *
 * What this file cannot do, at any price: check that what the model wrote is
 * what the document said. Nothing in this plumbing can. See the test.
 *
 * Nor does it check the language. The checkpoint named below declares
 * `language: en` on its card and was fine-tuned on English newswire; on a
 * French note it will write something, and nothing here will stop it. A
 * multilingual summariser exists — mT5 fine-tuned on XL-Sum, forty-five
 * languages, French among them — under CC BY-SA-NC 4.0, which rules out
 * commercial use. Read the card before you pick, and read it for the language
 * first.
 */

// A full stop, question or exclamation mark followed by whitespace, or a line
// break: a transcript without final punctuation is still cut into lines.
const SENTENCE_END = /(?<=[.!?])\s+|\s*\n\s*/;

// The ONNX conversion of sshleifer/distilbart-cnn-12-6, the weights Python
// loads. English only: that card says `language: en`, and the model is
// distilled from one fine-tuned on CNN/DailyMail and XSum, English newswire.
export const MODEL_NAME = 'Xenova/distilbart-cnn-12-6';

// What one call to the model is given to read. The window is 1,024 tokens, and
// characters are not tokens: with this model's tokenizer, 3,000 characters of
// French prose come to 1,091 tokens, 2,000 to 727. A piece that still
// overflows, code or a table, is refused by the model below, not truncated.
export const CHUNK_CHARACTERS = 2000;

// What the whole function is allowed to read. A self-hosted model costs
// machine time rather than money, but a document nobody meant to send is still
// better refused than churned through in silence.
export const MAX_CHARACTERS = 200_000;

export class SummaryUnavailable extends Error {}

/** The real model, loaded from local weights and run on this machine. */
export class LocalSummariser {
  #pipeline;

  async load(name = MODEL_NAME) {
    // Loads weights, so this is never reached in the tests.
    const { pipeline } = await import('@huggingface/transformers');
    this.#pipeline = await pipeline('summarization', name);
    return this;
  }

  async generate(text) {
    const { tokenizer } = this.#pipeline;
    // Refused rather than truncated: a summary of the start of a piece reads
    // exactly like a summary of all of it.
    if (tokenizer.encode(text).length > tokenizer.model_max_length) {
      throw new RangeError("this piece is longer than the model's window");
    }
    const [{ summary_text: written }] = await this.#pipeline(text);
    return written;
  }
}

let loading; // the default model, loaded once for the life of the process

function defaultModel() {
  loading ??= new LocalSummariser().load().catch((error) => {
    loading = undefined; // a failed load is tried again next time, not kept
    throw error;
  });
  return loading;
}

/** Join consecutive units into pieces of at most `size` characters. */
export function pack(units, size = CHUNK_CHARACTERS) {
  const pieces = [];
  let current = '';
  for (const unit of units) {
    if (current && current.length + 1 + unit.length > size) {
      pieces.push(current);
      current = unit;
    } else {
      current = `${current} ${unit}`.trim();
    }
  }
  if (current) pieces.push(current);
  return pieces;
}

/**
 * Cut the document into pieces that fit the model's window, at sentence
 * boundaries.
 *
 * A sentence longer than the window on its own is passed whole. Cutting
 * mid-sentence would hand the model half a clause, which is a worse thing to
 * summarise.
 */
export function chunk(text, size = CHUNK_CHARACTERS) {
  const sentences = text.trim().split(SENTENCE_END).map((part) => part.trim());
  return pack(sentences.filter(Boolean), size);
}

/**
 * @param {string} text
 * @param {object} options
 * @param {{generate: Function}} [options.model] injected so this can be tested
 *   without loading weights; defaults to the real model
 * @param {number} [options.attempts]
 */
export async function summarise(text, { model, attempts = 2 } = {}) {
  if (text.length > MAX_CHARACTERS) {
    throw new RangeError(`document longer than ${MAX_CHARACTERS} characters`);
  }
  let pieces = chunk(text);
  if (pieces.length === 0) return '';
  model ??= await defaultModel();

  for (;;) {
    const notes = [];
    for (const piece of pieces) notes.push(await generate(model, piece, attempts));
    if (notes.length === 1) return notes[0];
    // Next pass: the model reads back its own notes, packed to the window like
    // the document was. It never reads the document again.
    const packed = pack(notes);
    if (packed.length >= pieces.length) {
      throw new SummaryUnavailable('the notes are no shorter than what they summarise');
    }
    pieces = packed;
  }
}

async function generate(model, text, attempts) {
  let lastError;
  for (let i = 0; i < attempts; i += 1) {
    try {
      const written = ((await model.generate(text)) ?? '').trim();
      if (written) return written;
      // An empty answer is a failure, not a summary. Returning it would leave
      // a hole in the middle of the document with no trace.
      lastError = new Error('the model returned an empty summary');
    } catch (error) {
      lastError = error;
    }
  }
  throw new SummaryUnavailable(String(lastError));
}