Predecir el siguiente token en Java

Problema

Un comienzo, tres respuestas

El bot de soporte tiene que continuar respuestas. Un compañero ejecuta el mismo comienzo, «el pedido», tres veces con una semilla distinta para el generador. Estas son las tres respuestas:

var table = countNext(TEXT);
for (int seed : new int[] {2026, 34567, 12345})
    System.out.println("semilla " + seed + ": " + generate(table, "el pedido", 1, 100, new Rng(seed)));

semilla 2026: el pedido llegará en tres días
semilla 34567: el pedido llegó a un mensajero
semilla 12345: el pedido llegará en un plazo de 14 días

La primera suena razonable. La segunda parece casi una frase, pero el pedido «llegó a un mensajero» no tiene sentido. La tercera suena bien y es más larga. En las tres, cada par de palabras vecinas sale del texto de la tienda. Y las tres vienen del mismo programa, que no busca la respuesta en ningún sitio: escribe palabra por palabra.

¿Cómo sabe qué palabra poner? ¿Y por qué es distinto cada vez?

La idea

La probabilidad de la siguiente palabra

Un modelo de lenguaje recibe un texto y da una probabilidad para cada posible token siguiente. Nada más: no sabe qué es un pedido, solo sabe qué suele venir después de qué.

El texto se escribe con un bucle:

El modelo da las probabilidades del siguiente token.
Se elige uno según ellas.
Se añade al texto y todo se repite, hasta el final o el límite.

La elección es un paso aparte, no parte del modelo. Por eso el mismo modelo puede dar respuestas distintas.

El modelo de esta lección es el más pequeño posible: una tabla de cuántas veces vino cada palabra justo después de cada palabra. Mira solo la última palabra y cuenta palabras enteras, no tokens. Los modelos reales son redes neuronales que miran todo el texto, pero dan lo mismo: las probabilidades del siguiente token.

Míralo

La tabla: qué palabra después de cuál

El texto de la tienda es el mismo que en la lección sobre los tokens. Para cada frase contamos cada par de palabras vecinas. El punto «.» marca el final de la frase, para que el modelo sepa cuándo parar.

// una posible palabra siguiente y su probabilidad
record Choice(String word, double p) {}

// cuántas veces viene cada palabra justo después de cada palabra; "." termina una frase
static Map<String, Map<String, Integer>> countNext(String[] text) {
    Map<String, Map<String, Integer>> table = new LinkedHashMap<>();
    for (String line : text) {
        String[] w = (line + " .").split(" ");
        for (int i = 0; i + 1 < w.length; i++)
            table.computeIfAbsent(w[i], key -> new LinkedHashMap<>())
                 .merge(w[i + 1], 1, Integer::sum);
    }
    return table;
}

La probabilidad de la siguiente palabra es su conteo dividido entre la suma. Por ahora mira solo t = 1 y «todas las palabras»: entonces el peso es el conteo mismo. La temperatura t y k vienen más abajo.

// las k más frecuentes, peso conteo^(1/t) / suma
static List<Choice> nextWords(Map<String, Integer> counts,
                              double t, int k) {
    List<String> words = new ArrayList<>(counts.keySet());
    words.sort((a, b) -> counts.get(b) - counts.get(a));
    List<String> top =
        words.subList(0, Math.min(k, words.size()));
    double total = 0;
    for (String w : top)
        total += Math.pow(counts.get(w), 1 / t);
    List<Choice> choices = new ArrayList<>();
    for (String w : top) {
        double weight = Math.pow(counts.get(w), 1 / t);
        choices.add(new Choice(w, weight / total));
    }
    return choices;
}

var table = countNext(TEXT);
Map<String, Integer> counts = table.get("a");
System.out.println("después de \"a\":");
for (Choice c : nextWords(counts, 1, 100))
    System.out.printf("%-11s%2d  %5.1f%%%n", c.word(), counts.get(c.word()), c.p() * 100);

después de "a":
tu          3   42.9%
una         2   28.6%
la          1   14.3%
un          1   14.3%

Después de «a» el texto tiene cuatro palabras distintas, 7 veces en total. «tu» es 3 de ellas: 42,9 %. Eso es el modelo de lenguaje: para «a» dice «tu» con 42,9 %, «una» con 28,6 % y cada una de las otras dos, «la» y «un», con 14,3 %.

Míralo

La elección: un número de 0 a 1

¿Cómo se elige una palabra con probabilidad 42,9 %? El generador con semilla da un número r de 0 a 1. Las barras de las palabras se colocan una al lado de otra: «tu» de 0 a 0,43 aproximadamente, «una» de 0,43 a 0,71 y así sucesivamente. La palabra en cuya barra cae r es la elegida.

class Rng {
    private long state;

    Rng(long seed) {
        state = seed;
    }

    // el siguiente número entre 0 y 1; la misma semilla da los mismos números
    double next() {
        state = state * 48271 % 2147483647;
        return (double) state / 2147483647;
    }
}

// un número de 0 a 1 cae en una de las barras
static String pick(List<Choice> choices, Rng rng) {
    double r = rng.next();
    double sum = 0;
    for (Choice c : choices) {
        sum += c.p();
        if (r < sum)
            return c.word();
    }
    return choices.get(choices.size() - 1).word();
}

La generación repite esto palabra por palabra desde el comienzo, hasta que llega «.» o se alcanzan 12 palabras:

// continúa el texto palabra por palabra hasta "." o 12 palabras
static String generate(Map<String, Map<String, Integer>> table,
                       String start, double t, int k, Rng rng) {
    List<String> text = new ArrayList<>(List.of(start.split(" ")));
    String word = text.get(text.size() - 1);
    while (text.size() < 12) {
        word = pick(nextWords(table.get(word), t, k), rng);
        if (word.equals("."))
            break;
        text.add(word);
    }
    return String.join(" ", text);
}

Si siempre tomas la palabra más probable (k = 1), el generador no tiene nada que elegir y la respuesta es la misma con cualquier semilla:

var table = countNext(TEXT);
for (int seed : new int[] {2026, 34567, 12345})
    System.out.println("semilla " + seed + ": " + generate(table, "el pedido", 1, 1, new Rng(seed)));

semilla 2026: el pedido llegará en tres días
semilla 34567: el pedido llegará en tres días
semilla 12345: el pedido llegará en tres días

Con elección aleatoria la respuesta depende de la semilla. Pero la misma semilla da los mismos números, y por tanto la misma respuesta:

var table = countNext(TEXT);
for (int seed : new int[] {2026, 2026, 34567})
    System.out.println("semilla " + seed + ": " + generate(table, "el pedido", 1, 100, new Rng(seed)));

semilla 2026: el pedido llegará en tres días
semilla 2026: el pedido llegará en tres días
semilla 34567: el pedido llegó a un mensajero

Por eso las tres respuestas del comienzo eran distintas: cada semilla da otros números y, con ellos, otras palabras. La semilla 2026 coincide con la elección más probable, pero las otras dos se desvían de ella.

Pruébalo

Predice la salida

Después de «entrega» el texto tiene cuatro continuaciones, una de ellas el final «.». ¿Qué imprime el programa con k = 1?

var table = countNext(TEXT);
for (Choice c : nextWords(table.get("entrega"), 1, 1))
    System.out.printf("%s %.1f%%%n", c.word(), c.p() * 100);

Respuesta

a 40.0%

a 40.0%
. 20.0%
es 20.0%
tarda 20.0%

a 100.0%

La idea

Temperatura y top-k

La elección tiene dos mandos, que también tienen los modelos reales:

temperatura t: el peso es conteo^(1/t). Con t menor que 1 las palabras frecuentes se vuelven aún más probables y la respuesta es más predecible. Con t mayor que 1 las probabilidades se igualan y las palabras raras salen más a menudo.
top-k: solo quedan las k palabras más frecuentes y las demás no se pueden elegir. La suma se calcula solo sobre ellas, por eso vuelve a ser 100 %.

var table = countNext(TEXT);
Map<String, Integer> counts = table.get("a");
double[] temps = {0.5, 1.0, 2.0};
List<List<Choice>> rows = new ArrayList<>();
StringBuilder head = new StringBuilder(" ".repeat(11));
for (double t : temps) {
    rows.add(nextWords(counts, t, 100));
    head.append(String.format("%7s", "t=" + (t == (int) t ? String.valueOf((int) t) : String.valueOf(t))));
}
System.out.println(head);
for (int i = 0; i < counts.size(); i++) {
    StringBuilder line = new StringBuilder(String.format("%-11s", rows.get(0).get(i).word()));
    for (List<Choice> row : rows)
        line.append(String.format("%6.1f%%", row.get(i).p() * 100));
    System.out.println(line);
}

             t=0.5    t=1    t=2
tu           60.0%  42.9%  33.7%
una          26.7%  28.6%  27.5%
la            6.7%  14.3%  19.4%
un            6.7%  14.3%  19.4%

Con t = 0,5 «tu» es 60,0 %, y con t = 2 solo 33,7 %. Las palabras vistas una sola vez («la» y «un») tienen 6,7 %, 14,3 % y 19,4 %: cuanto más alta es la temperatura, más a menudo las elige el modelo.

Aquí están las barras después de «a» con las tres temperaturas, luego con k = 3 y la elección con un número del generador:

// las k más frecuentes, peso conteo^(1/t) / suma
static List<Choice> nextWords(Map<String, Integer> counts,
                              double t, int k) {
    List<String> words = new ArrayList<>(counts.keySet());
    words.sort((a, b) -> counts.get(b) - counts.get(a));
    List<String> top =
        words.subList(0, Math.min(k, words.size()));
    double total = 0;
    for (String w : top)
        total += Math.pow(counts.get(w), 1 / t);
    List<Choice> choices = new ArrayList<>();
    for (String w : top) {
        double weight = Math.pow(counts.get(w), 1 / t);
        choices.add(new Choice(w, weight / total));
    }
    return choices;
}

// un número de 0 a 1 cae en una de las barras
static String pick(List<Choice> choices, Rng rng) {
    double r = rng.next();
    double sum = 0;
    for (Choice c : choices) {
        sum += c.p();
        if (r < sum)
            return c.word();
    }
    return choices.get(choices.size() - 1).word();
}

Paso a paso, en palabras

Después de «a» el texto tiene 4 palabras. Una barra es el conteo dividido entre los 7: «tu» es 3 de 7, 42,9 %.
t = 0,5: el peso es conteo². La palabra frecuente se vuelve aún más probable: «tu» 60,0 %.
t = 2: el peso es √conteo. Las barras se igualan: «tu» 33,7 %, «un» 19,4 %.
t = 1, k = 3: solo quedan las tres más frecuentes, y juntas vuelven a sumar 100 %: 50,0 %, 33,3 %, 16,7 %.
Las tres barras quedan una junto a otra de 0 a 1. El generador con la semilla 31337 da 0,7044.
0,7044 cae en la barra de «una» (de 50,0 % a 83,3 %): la siguiente palabra es «una», no la más frecuente.

var table = countNext(TEXT);
List<Choice> choices = nextWords(table.get("a"), 1, 3);
System.out.println("k=3:");
for (Choice c : choices)
    System.out.printf("%-11s%5.1f%%%n", c.word(), c.p() * 100);
Rng rng = new Rng(31337);
Rng copy = new Rng(31337);
System.out.printf("número del generador: %.4f%n", copy.next());
System.out.println("elegida: " + pick(choices, rng));

k=3:
tu          50.0%
una         33.3%
la          16.7%
número del generador: 0.7044
elegida: una

Con k = 3 quedan «tu», «una» y «la», y sus conteos son 3, 2 y 1 de 6. El generador con la semilla 31337 da 0,7044, y eso cae en la barra de «una», de 50,0 % a 83,3 %. La elección sigue las probabilidades, pero no siempre es la palabra más probable.

Pruébalo

Completa

¿Qué va en lugar de ___ en el peso de una palabra?

for (String w : top)
    total += Math.pow(counts.get(w), 1 / );

Pruébalo

Encuentra el error

Un compañero quiere top-k con las tres palabras más probables, pero su modelo escribe solo palabras raras. ¿En qué línea está el error?

Elige la línea con el error.

La idea

Cuándo duele

Lo probable no es lo verdadero. El modelo escribe «llegará en tres días» con la misma facilidad con que escribe «tarda hasta dos días»: las dos están en el texto. No sabe cuándo llega este pedido en concreto. El texto fluido suena seguro aunque sea incorrecto. De eso trata la lección sobre los límites del modelo.
Cada par es correcto, la frase no. «el pedido llegó a un mensajero» está hecha de palabras vecinas que existen en el texto. Nuestro modelo mira solo una palabra atrás. Los reales miran todo el texto y escriben mucho mejor, pero aun así pueden pegar algo fluido y sin sentido.
Más temperatura, más errores. Las continuaciones raras salen más a menudo. Para respuestas a clientes, donde importa la exactitud, la temperatura es baja. Para ideas y variantes puede ser más alta.
Una semilla no es una promesa. En la lección una semilla da una respuesta. Con un modelo real detrás de una API, la misma solicitud a veces da otra respuesta incluso con los mismos ajustes, así que no cuentes con obtener exactamente el mismo texto: comprueba la respuesta.

Más precisamente, para curiosos

Softmax y temperatura. Un modelo real es una red neuronal: para cada token del vocabulario da un número (un logit), y softmax los convierte en probabilidades. La temperatura divide los números entre t antes de softmax. El peso conteo^(1/t) de la lección es justo eso: softmax de log(conteo) / t.

top-p. Además de top-k existe top-p: se quedan los tokens más frecuentes hasta que su suma llega a p (por ejemplo 0,9). Así, cuando el modelo está seguro quedan pocos tokens, y cuando duda, más.

Tokens, no palabras. Un modelo real elige tokens de su vocabulario, y una palabra sale de varias elecciones seguidas. La temperatura, top-k y la semilla se fijan en la solicitud al modelo (la lección sobre la llamada).

Cuánto mira atrás. Nuestro modelo recuerda una palabra. Uno real ve toda la conversación hasta su límite, y eso es la ventana de contexto.

Ejercicio

Qué tan probable es la frase

Resuélvelo en tu editor.

La probabilidad de una frase entera es el producto de las probabilidades de cada palabra siguiente, incluido el «.» del final. Escribe una función que la calcule con la tabla de la lección e imprímela en porcentaje para tres frases. Si un par de palabras nunca se ha visto, su probabilidad es 0.

el paquete llegará en tres días: 1.02%
el paquete llegará en su embalaje original: 0.51%
el paquete llegó mañana: 0.00%

La primera frase tiene 1,02 %, aunque está en el texto palabra por palabra: en casi cada paso el modelo habría podido elegir otra palabra. La segunda tiene 0,51 %; no está en el texto, pero todos sus pares se han visto. La tercera tiene 0,00 %: «paquete» y «llegó» están las dos en el texto, pero nunca una justo después de la otra. Así dice el modelo qué tan «natural» le parece un texto.

Más tareas

La frase más probable. Empieza en «el paquete» y toma siempre la palabra más probable. ¿Cuál es la probabilidad de la frase que sale?
Dos palabras atrás. Cuenta qué palabra viene después de un par de palabras, no de una. ¿Cómo cambian las respuestas del comienzo de la lección?

Test

Ponte a prueba

1. ¿Qué te da un modelo de lenguaje cuando le pasas un texto?Solo una palabra: la más probableLa respuesta terminada, encontrada en el texto con el que aprendióUna probabilidad para cada posible palabra siguiente (token)

2. Ejecutas la generación dos veces con la semilla 2026. ¿Qué ves?Dos textos distintos, porque la elección es aleatoriaEl mismo texto: una semilla da los mismos númerosUn error: una semilla solo se puede usar una vez

3. ¿Cómo cambia la temperatura t = 0,5 las probabilidades después de «a»?No cambia nada, solo la velocidadLa palabra frecuente se vuelve aún más probable y las raras aún menosTodas las palabras se vuelven igual de probables

4. El bot escribe «el pedido llegará en tres días», pero la entrega en realidad tarda hasta dos días. ¿Por qué?El modelo escribe una continuación probable, no un hecho comprobadoLa temperatura es demasiado bajaEl generador está roto

Predecir el siguiente token

Cookies