110 lines
3.6 KiB
Python
110 lines
3.6 KiB
Python
# Ported from ref/word_cloud/wordcloud/tokenization.py (MIT License)
|
|
from __future__ import division
|
|
from itertools import tee
|
|
from operator import itemgetter
|
|
from collections import defaultdict
|
|
from math import log
|
|
|
|
|
|
def _l(k, n, x):
|
|
"""Dunning log-likelihood helper."""
|
|
return log(max(x, 1e-10)) * k + log(max(1 - x, 1e-10)) * (n - k)
|
|
|
|
|
|
def _collocation_score(count_bigram, count1, count2, n_words):
|
|
"""Dunning likelihood ratio collocation score."""
|
|
if n_words <= count1 or n_words <= count2:
|
|
return 0
|
|
N, c12, c1, c2 = n_words, count_bigram, count1, count2
|
|
p = c2 / N
|
|
p1 = c12 / c1
|
|
p2 = (c2 - c12) / (N - c1)
|
|
score = (
|
|
_l(c12, c1, p) + _l(c2 - c12, N - c1, p)
|
|
- _l(c12, c1, p1) - _l(c2 - c12, N - c1, p2)
|
|
)
|
|
return -2 * score
|
|
|
|
|
|
def _pairwise(iterable):
|
|
a, b = tee(iterable)
|
|
next(b, None)
|
|
return zip(a, b)
|
|
|
|
|
|
def process_tokens(words, normalize_plurals=True):
|
|
"""Count words, normalizing case and optionally merging plurals.
|
|
|
|
Returns
|
|
-------
|
|
counts : dict str -> int
|
|
standard_forms : dict lowercase_str -> canonical_str
|
|
"""
|
|
d = defaultdict(dict)
|
|
for word in words:
|
|
wl = word.lower()
|
|
case_dict = d[wl]
|
|
case_dict[word] = case_dict.get(word, 0) + 1
|
|
|
|
if normalize_plurals:
|
|
merged_plurals = {}
|
|
for key in list(d.keys()):
|
|
if key.endswith('s') and not key.endswith('ss'):
|
|
singular = key[:-1]
|
|
if singular in d:
|
|
for word, count in d[key].items():
|
|
sing_form = word[:-1]
|
|
d[singular][sing_form] = d[singular].get(sing_form, 0) + count
|
|
merged_plurals[key] = singular
|
|
del d[key]
|
|
|
|
fused_cases = {}
|
|
standard_cases = {}
|
|
item1 = itemgetter(1)
|
|
for word_lower, case_dict in d.items():
|
|
first = max(case_dict.items(), key=item1)[0]
|
|
fused_cases[first] = sum(case_dict.values())
|
|
standard_cases[word_lower] = first
|
|
|
|
if normalize_plurals:
|
|
for plural, singular in merged_plurals.items():
|
|
standard_cases[plural] = standard_cases.get(singular, singular)
|
|
|
|
return fused_cases, standard_cases
|
|
|
|
|
|
def unigrams_and_bigrams(words, stopwords, normalize_plurals=True,
|
|
collocation_threshold=30):
|
|
"""Return word counts including statistically significant bigrams."""
|
|
bigrams = [
|
|
p for p in _pairwise(words)
|
|
if not any(w.lower() in stopwords for w in p)
|
|
]
|
|
unigrams = [w for w in words if w.lower() not in stopwords]
|
|
n_words = len(unigrams)
|
|
|
|
counts_unigrams, standard_form = process_tokens(
|
|
unigrams, normalize_plurals=normalize_plurals)
|
|
counts_bigrams, _ = process_tokens(
|
|
[" ".join(b) for b in bigrams], normalize_plurals=normalize_plurals)
|
|
|
|
orig_counts = counts_unigrams.copy()
|
|
|
|
for bigram_string, count in counts_bigrams.items():
|
|
parts = bigram_string.split(" ", 1)
|
|
if len(parts) != 2:
|
|
continue
|
|
word1 = standard_form.get(parts[0].lower(), parts[0])
|
|
word2 = standard_form.get(parts[1].lower(), parts[1])
|
|
if word1 not in orig_counts or word2 not in orig_counts:
|
|
continue
|
|
score = _collocation_score(count, orig_counts[word1],
|
|
orig_counts[word2], n_words)
|
|
if score > collocation_threshold:
|
|
counts_unigrams[word1] -= count
|
|
counts_unigrams[word2] -= count
|
|
counts_unigrams[bigram_string] = count
|
|
|
|
# Remove non-positive counts
|
|
return {w: c for w, c in counts_unigrams.items() if c > 0}
|