diora-web/feeds/cluster.py

"""Bundling news items about the same event into stories.

Each item is described by the words of its title (counted twice) and summary,
without stop words and with a crude German/English suffix cut. A new item joins
the most similar story of the last WINDOW hours (cosine of tf-idf vectors,
document frequencies over those stories) if that similarity reaches THRESHOLD
and they share at least MIN_SHARED words; otherwise it starts a story of its own.
A story many sources report on is a top story.
"""

import math
import re
from collections import Counter
from datetime import timedelta

from .models import Item, Story

WINDOW = timedelta(hours=36)
THRESHOLD = 0.25
MIN_SHARED = 2

STOP = set('''
aber alle allem allen aller alles also als am an andere anderen auch auf aus bei beim bereits bis bisher
da dabei damit dann darauf darf darum das dass dem den denen der deren des deshalb dessen die dies diese
diesem diesen dieser dieses doch dort durch ein eine einem einen einer eines einige er es etwa euch für
gegen geht gibt hab habe haben hat hatte hatten hier hin hinter ich ihr ihre ihrem ihren ihrer im immer
in ins ist ja jahr jahre jahren jetzt kann kein keine können könnte man mehr mit muss müssen nach neue neuen
nicht nichts noch nun nur ob oder ohne schon sei seid sein seine seinem seinen seiner seit sich sie sind
so soll sollen sollte sondern sowie über um und uns unter viel viele vom von vor war waren was weil weiter
weitere wenn wer werde werden wie wieder will wir wird wo worden wurde wurden zu zum zur zwei drei vier
heute gestern morgen laut sagt sagte erst mal neu alt video live ticker liveblog news meldung bericht
the and for with from that this have has are was were will would been into about after over their they
his her its our you your not but all can new says said more than when what who how why off out
tot tote toten getötet sterben stirbt gestorben menschen mensch leute millionen milliarden prozent euro
deutschland deutsche deutschen land länder stadt welt woche wochen monat monate tag tage tagen uhr
frau frauen mann männer kinder erste ersten zweite neuer neues großen große groß kleine wegen gibt geben
'''.split())

_WORD = re.compile(r'[a-zäöüß][a-zäöüß0-9\-]+')
_SUFFIXES = ('ungen', 'ung', 'en', 'er', 'es', 'e', 's', 'n')


def _stem(w):
    if len(w) > 5:
        for s in _SUFFIXES:
            if w.endswith(s) and len(w) - len(s) >= 4:
                return w[:-len(s)]
    return w


def words(text):
    out = []
    for w in _WORD.findall((text or '').lower()):
        w = w.strip('-')
        if len(w) < 3 or w in STOP:
            continue
        out.append(_stem(w))
    return out


def item_terms(title, summary):
    return Counter(words(title) * 2 + words(summary[:300]))


def _cosine(a, b, idf):
    dot = sum(a[w] * b[w] * idf.get(w, 1.0) ** 2 for w in a.keys() & b.keys())
    if not dot:
        return 0.0
    na = math.sqrt(sum((n * idf.get(w, 1.0)) ** 2 for w, n in a.items()))
    nb = math.sqrt(sum((n * idf.get(w, 1.0)) ** 2 for w, n in b.items()))
    return dot / (na * nb)


class Clusterer:
    """Assigns one user's new news items to stories (load once, add items, done).

    An item is compared with every item already in a story: the story's score
    is the mean of the best match and the average match, so one stray shared
    word can't pull an item in, and a story doesn't grow into a magnet the way
    a pooled word list would. Items of the same source join only when very
    close (a source rarely reports one event twice, but often two similar ones)."""

    SAME_SOURCE = 0.5

    def __init__(self, user, now):
        self.user = user
        self.stories = list(Story.objects.filter(user=user, last_at__gte=now - WINDOW))
        self.members = {s.pk: [] for s in self.stories}   # story id -> [(source id, terms)]
        for i in Item.objects.filter(story__in=self.stories, gone=False).only('story_id', 'source_id', 'title', 'summary'):
            self.members[i.story_id].append((i.source_id, item_terms(i.title, i.summary)))
        self._df = Counter()
        for ms in self.members.values():
            for _, t in ms:
                self._df.update(t.keys())
        self._n = sum(len(ms) for ms in self.members.values())

    def _idf(self):
        n = self._n + 1
        return {w: math.log((n + 1) / (d + 1)) + 1 for w, d in self._df.items()}

    def score(self, t, source_id, members, idf):
        sims = []
        for sid, mt in members:
            if len(t.keys() & mt.keys()) < MIN_SHARED:
                sims.append(0.0)
                continue
            sim = _cosine(t, mt, idf)
            if sid == source_id and sim < self.SAME_SOURCE:
                sim = 0.0
            sims.append(sim)
        if not sims or max(sims) == 0:
            return 0.0
        return (max(sims) + sum(sims) / len(sims)) / 2

    def add(self, item):
        t = item_terms(item.title, item.summary)
        best, best_sim = None, 0.0
        if t:
            idf = self._idf()
            for s in self.stories:
                if abs((s.last_at - item.published).total_seconds()) > WINDOW.total_seconds():
                    continue
                sim = self.score(t, item.source_id, self.members[s.pk], idf)
                if sim > best_sim:
                    best, best_sim = s, sim
        if best is not None and best_sim >= THRESHOLD:
            self._join(best, item)
        else:
            best = Story.objects.create(user=self.user, lead=item, first_at=item.published,
                                        last_at=item.published, size=1, sources=1)
            self.stories.append(best)
            self.members[best.pk] = []
            item.story = best
            item.save(update_fields=['story'])
        self.members[best.pk].append((item.source_id, t))
        self._df.update(t.keys())
        self._n += 1
        return best

    def _join(self, s, item):
        item.story = s
        item.save(update_fields=['story'])
        s.size = Item.objects.filter(story=s).count()
        s.sources = Item.objects.filter(story=s).values('source').distinct().count()
        s.first_at = min(s.first_at, item.published)
        s.last_at = max(s.last_at, item.published)
        if s.lead is None or (not s.lead.image_url and item.image_url):
            s.lead = item
        s.save(update_fields=['size', 'sources', 'first_at', 'last_at', 'lead'])