Extract TF-IDF Keywords from Text Documents

from __future__ import annotations

import math
import re
from collections import Counter

TOKEN_RE = re.compile(r"[a-z][a-z0-9+#.-]{1,}")
STOPWORDS = {
    "the", "and", "for", "with", "that", "this",
    "from", "into", "your", "are", "was", "have",
}


def tokenize(text: str) -> list[str]:
    return [
        token
        for token in TOKEN_RE.findall(text.lower())
        if token not in STOPWORDS and len(token) > 2
    ]


def keyword_scores(documents: list[str], top_n: int = 5) -> list[tuple[str, float]]:
    tokenized = [tokenize(document) for document in documents]
    document_frequency: Counter[str] = Counter()
    for tokens in tokenized:
        document_frequency.update(set(tokens))
    total_documents = len(tokenized)
    scores: Counter[str] = Counter()
    for tokens in tokenized:
        counts = Counter(tokens)
        length = sum(counts.values()) or 1
        for term, count in counts.items():
            tf = count / length
            idf = math.log((1 + total_documents) / (1 + document_frequency[term])) + 1
            scores[term] += tf * idf
    return sorted(scores.items(), key=lambda item: (-item[1], item[0]))[:top_n]


def main() -> None:
    documents = [
        "Use connection pooling to reduce database latency and reuse open sockets.",
        "A connection pool keeps sockets warm so database queries avoid handshake latency.",
        "Caching responses can hide database latency but risks stale data.",
    ]
    for term, score in keyword_scores(documents, 4):
        print(f"{term:12s} {score:.4f}")
    assert keyword_scores(documents, 1)[0][0] == "database"


if __name__ == "__main__":
    main()