alexeygrigorev/data-science-interviews: PMI пар слов

Условие

По корпусу текстов посчитайте PMI (pointwise mutual information) для всех биграмм слов: PMI(x, y) = log( P(x, y) / ( P(x) * P(y) ) ), где P(x,y) — вероятность встретить пару (x, y) рядом, P(x) — частота слова x.

Решение

Подход

Считаем три счётчика: n_x (одиночные слова), n_xy (пары соседних слов), N (общее число слов). Затем P(x) = n_x / N, P(x, y) = n_xy / (N-1) и PMI = log(P(xy)/(P(x)*P(y))).

Реализация

from collections import Counter
from math import log
 
def pmi(docs: list[list[str]]) -> dict[tuple[str, str], float]:
    unigram = Counter()
    bigram = Counter()
    total_uni = 0
    total_bi = 0
    for doc in docs:
        unigram.update(doc)
        total_uni += len(doc)
        for a, b in zip(doc, doc[1:]):
            bigram[(a, b)] += 1
            total_bi += 1
    result = {}
    for (a, b), n_ab in bigram.items():
        p_ab = n_ab / total_bi
        p_a = unigram[a] / total_uni
        p_b = unigram[b] / total_uni
        result[(a, b)] = log(p_ab / (p_a * p_b))
    return result

Подводные камни

1. PMI завышает редкие пары: пара, встречающаяся ровно раз, может дать большое значение. Поэтому на практике используют PMI с фильтрацией по n_ab >= threshold или Positive PMI (max(0, pmi)).
2. Если пара ни разу не встретилась — PMI неопределён; считаем только наблюдаемые пары.
3. Контекст «рядом» неоднозначен: окно ±1, ±2 или предложение? Уточняем заранее.

Эталонный ответ

Через Counter для униграмм и биграмм; нормируем на суммы и считаем log(P(xy)/(P(x)P(y))). Для прода — PPMI с фильтром по частоте.