Chinese Placement Test with Logistic Regression and Active Learning

import re
import time
import random
from datetime import datetime
from itertools import combinations
from collections import defaultdict
from textwrap import wrap

import jieba
import numpy as np
import seaborn as sns
import pandas as pd
import scipy.stats as st
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
import ipywidgets as widgets
from IPython.display import display

%matplotlib inline

#-----------------------------------------------------------------------
# Load the Cedict dictionary
#-----------------------------------------------------------------------
dictionary = defaultdict(list)
filter_out_patterns = [
    'see ',
    'abbr. for ',
    'see also ',
    'also written ',
    'euphemistic variant of',
    'variant of ',
    'unofficial variant of ',
    'archaic variant of ',
    'old variant of ',
    'ancient variant of ',
    'erhua variant of ',
    'Japanese variant of ',
    'CL:',
    '(old)',
    '(dialect)',
    '(surname)',
    '(Cantonese)',
    'radical',
    'Mandarin equivalent'
]

with open('cedict_ts.u8') as f:
    for i, line in enumerate(f):
        if line.startswith('#'): # Skip comments
            continue

        _, hz, py, transl = re.match(
            r"(\S*) (\S*) \[(.*)\] \/(.*)\/", line).groups()
        
        if py[0].isupper():
            continue # Skip named entities

        t_split = transl.split('/')
        t_filtered = []
        for t in t_split:
            if t.startswith('surname '):
                continue
            not_note_pattern = r'[\w\W]*\([\w\W]* pr\. [\w\W]*\)'
            if ' pr. ' in t and not re.match(not_note_pattern, t):
                # short for pronounciation. There's:
                # also pr. , Taiwain pr. , Japan pr.
                # However, when the pr. is inside parenthesis, it applies to
                # the translation which comes before
                continue

            filter_out = False
            for pattern in filter_out_patterns:
                if pattern in t:
                    filter_out = True
            if filter_out:
                continue
            t_filtered.append(t)
                    
        if len(t_filtered) > 0:
            dictionary[hz].append((py, t_filtered))

#-----------------------------------------------------------------------
# Load Weibo word frequencies
#-----------------------------------------------------------------------
word_freq = {}
word_freq_all = {}
with open('weibo_wordfreq.release_UTF-8.txt', encoding='utf-8-sig') as f: 
    for line in f.readlines():
        word, count = line.split('\t')
        if word in word_freq:
            continue

        word_freq_all[word] = int(count)
        if len(dictionary[word]) == 0:
            continue

        word_freq[word] = int(count)

words = list(word_freq.keys())
# Correct the frequency of components (component can not have lower frequency than compound)
for w in words:
    f = word_freq[w]
    if f < 50000: continue # skip uncommon words to speed up
    substrs = [w[x:y] for x, y in combinations(range(len(w)+1), r=2)]
    for substr in substrs:
        if substr not in words: continue
        subf = word_freq.get(substr, 0)
        if subf < f:
            word_freq[substr] = f
            
# Calculate word rank, higher for higher frequencies
word_rank = {w: r for r, (w, _) in 
             enumerate(sorted(word_freq.items(), key=lambda x: x[1]))}

#-----------------------------------------------------------------------
# Load the HSK word lists
#-----------------------------------------------------------------------
hsk_words_by_lvl = {}
hsk_lvl_by_word = {}
for lvl in range(1, 7):
    words = []
    with open(f'hsk/HSK{lvl}.txt', encoding='utf-8-sig') as f:
        lines = list(f.readlines())
        for w in lines:
            w = w.strip() # strip away newlines
            if w not in word_freq: continue # some HSK words are not in dict
            words.append((w, word_freq[w]))
            hsk_lvl_by_word[w] = lvl
    hsk_words_by_lvl[lvl] = words
    
#-----------------------------------------------------------------------
# Add components of HSK words as separate HSK words, e.g. 后面 -> 后， 面
# This is because the HSK lists are missing most of these components
#-----------------------------------------------------------------------
# First, build an index over all HSK words for all their substrings
hsk_word_index = {}
for lvl in range(1, 7):
    for w, f in hsk_words_by_lvl[lvl]:
        substrs = [w[x:y] for x, y in combinations(range(len(w)+1), r=2)]
        for substr in substrs:
            _, _, curr_lvl = hsk_word_index.get(substr, (None, None, 7))
            if lvl < curr_lvl:
                hsk_word_index[substr] = (w, f, lvl)

# Next, go through all non-HSK words and see if they are a component in an HSK word
# If so, add them to the corresponding HSK level (if also in dictionary)
component_assignments_per_lvl = defaultdict(list)
for w, f in word_freq.items():
    if w in hsk_lvl_by_word:
        continue # Already is HSK word

    if w in hsk_word_index:
        hsk_w, f, lvl = hsk_word_index[w]
        hsk_words_by_lvl[lvl].append((w, f))
        hsk_lvl_by_word[w] = lvl
        component_assignments_per_lvl[lvl].append((w, hsk_w))

# Some debug prints
if False:
    for lvl in range(1, 4):
        words = ' '.join([w for w, _ in component_assignments_per_lvl[lvl]])
        print(f'Component HSK {lvl} words:')
        print('\n'.join(wrap(words, 80)))

    print('\nAdded for each level: ')
    for lvl in range(1, 7):
        print(f'HSK{lvl}: {len(component_assignments_per_lvl[lvl])}')

freqs = np.array([f for words in hsk_words_by_lvl.values()
                  for _, f in words])
lvls = np.array([lvl for lvl, words in hsk_words_by_lvl.items()
                 for w in words])
fig, ax = plt.subplots(figsize=(10, 5))
sns.boxplot(x=lvls, y=freqs, ax=ax, fliersize=0);
ax.set(ylim=(0, 1.4e7), xlabel='HSK Level', ylabel='Frequency');

# Find the k most frequent words and their HSK level
K = 5000
common_k_words = sorted(list(word_freq.items()), key=lambda x: x[1])[-K:]
common_k_words_lvls = [(w, hsk_lvl_by_word.get(w, None)) for w, f in common_k_words]

color_map = {None: 'white', 1: 'blue', 2: 'orange', 3: 'green', 4: 'red', 5: 'purple', 6: 'brown'}
colors = [color_map[lvl] for _, lvl in common_k_words_lvls]
plt.figure(figsize=(15, 2))
ax = plt.axes([0,0,1,1], frameon=False)
ax.get_xaxis().set_visible(False)
ax.get_yaxis().set_visible(False)
plt.bar(np.arange(len(colors)), np.ones(len(colors)), width=1.0, color=colors)
plt.show()

# Reverse order to lowest frequency first
counts = np.array(list(word_freq_all.values()))
counts_subrange = counts[(counts > 1000) & (counts < 100000)]
plt.rcParams['figure.figsize'] = [15, 5] # set larger plots
plt.subplot(1, 2, 1)
plt.hist(counts_subrange, 100)
plt.ylabel('Count')
plt.xlabel('Word Frequency')
plt.subplot(1, 2, 2)
plt.hist(np.log(counts_subrange), 100, log=True)

plt.show()

word_rank_list = sorted(list(word_rank.items()), key=lambda x: x[1])
words = [w for w, _ in word_rank_list]
ranks = np.arange(len(word_rank_list))
in_hsk = np.array([w in hsk_word_index for w in words])
hsk_lvls = np.array([hsk_lvl_by_word.get(w, 0) for w in words])

# -----------------------------------------------------------------------
# Here are a few different versions of the hsk variable
# NOTE: out of these options, quadratic has the lowest mean loss and std
# -----------------------------------------------------------------------
# Unused:
#hsk_data = []
# Quadratic:
hsk_data = [np.array([(7 - hsk_lvl_by_word.get(w, 7))**2/36 for w in words])]
# Linear:
#hsk_data = [np.array([(7 - hsk_lvl_by_word.get(w, 7))/6 for w in words])]
# Single binary:
#hsk_data = [np.array([1.0 if w in hsk_lvl_by_word else 0.0 for w in words])]
# Per level binary:
#hsk_data = [np.array([1.0 if hsk_lvl_by_word.get(w, None) == lvl else 0.0 for w in words])
#            for lvl in range(1, 7)]

data = np.vstack((ranks.astype(float), *hsk_data)).T
data_hsk_mean = data[in_hsk, 0].mean(axis=0)
data_hsk_std = data[in_hsk, 0].std(axis=0)
normalized_data = data.copy()
normalized_data[:, 0] = (data[:, 0] - data_hsk_mean) / data_hsk_std
normalized_word_rank = dict(zip(words, normalized_data[:, 0]))

def sample_HSK_words_per_lvl(K):
    test_words = []
    random.seed(1) # Make sure we randomize reproducibly
    for lvl in range(1, 7):
        test_words = test_words + random.sample(hsk_words_by_lvl[lvl], K)
    return test_words

K = 50
test_words = sample_HSK_words_per_lvl(K)
test_rank = np.array([normalized_word_rank[w] for w, _ in test_words]).reshape(-1, 1)
# Set the lower level words as known, rest as unknown
outcomes = np.zeros(len(test_words))
outcomes[:len(test_words) // 2] = 1

def plot_logistic_regression(X, y, lr, show=True):
    plt.figure()
    plt.scatter(X.ravel(), y, color='black', zorder=20)
    X_test = np.linspace(-0.5, 1.25, 1000)
    y_plot = lr.predict_proba(X_test.reshape(-1, 1))[:, 1]
    plt.plot(X_test, y_plot, color='red', linewidth=3)
    plt.ylabel('Probability')
    plt.xlabel('Normalized Frequency rank')
    plt.xlim(-0.5, 1.25)
    if show: plt.show()

lr = LogisticRegression(random_state=0, solver='lbfgs', C=1e9)
lr.fit(test_rank, outcomes)
plot_logistic_regression(test_rank, outcomes, lr)

plt.figure()
X = np.linspace(-7.0, 7.0, 1000)
y = 1 / (1 + np.exp(-X))
E = - y * np.log(y) - (1 - y) * np.log(1-y)
plt.plot(X, y, color='red', linewidth=3)
plt.plot(X, E, color='blue', linewidth=2)
plt.ylabel('y')
plt.xlabel('X')
plt.figtext(.15, 0, "Red line: the logistic function, blue line: the entropy", fontsize=16)
plt.show()

HANZI_UTF_RANGES = [
    ('\u4E00', '\u9FFF'),
    ('\u3400', '\u4DBF'),
    ('\uF900', '\uFAFF')
]

def is_hanzi(char):
    for start, end in HANZI_UTF_RANGES:
        if ord(char) >= ord(start) and ord(char) <= ord(end):
            return True
    return False

def filter_hanzi(text):
    return ''.join(char for char in text if is_hanzi(char))

anki_words = set()
with open('anki_notes.txt') as f:
    for line in f:
        hanzi = filter_hanzi(line)
        #print(line)
        for character in hanzi:
            if character in words:
                anki_words.add(character)
        for token, *_ in jieba.tokenize(hanzi, mode='search'):
            if token in words:
                anki_words.add(token)
        if len(dictionary[hanzi]) != 0:
            anki_words.add(token)
                
# Add words from HSK levels I learned before using Anki
for lvl in range(1, 5):
    for w, _ in hsk_words_by_lvl[lvl]:
        anki_words.add(w)
        
# Add manually gathered extra words
with open('known_words.txt') as f:
    for line in f:
        anki_words.add(line.strip())

seed = datetime.now().microsecond
random.seed(seed)
print('Using seed:', seed)
np.seterr(all='raise')

component_words = set()
for lvl in range(1, 7):
    component_words.update(set(w for w, _ in component_assignments_per_lvl[lvl]))
    
def prevent_nans(probs):
    probs[probs[:, 0] == 0.0, 0] = 1-1e-20 # Prevent NaNs
    probs[probs[:, 1] == 1.0, 1] = 1e-20

def lr_prediction_and_sampling_cdf(lr, plot=False):
    probs = lr.predict_proba(normalized_data)
    low_prob = probs[:, 1] < 0.01
    prevent_nans(probs)
    try:
        expected_information = -probs*np.log(probs)
    except:
        breakpoint()
    expected_information = expected_information.sum(axis=1)
    lr_probs = expected_information / expected_information.sum()
    # Clip the long tail, to compensate for bias (very low probability words)
    #lr_probs[low_prob] = 0
    lr_cdf = np.cumsum(lr_probs)
    if plot:
        plt.plot(normalized_data[::100, 0], lr_probs[::100])
        plt.show()
    return probs, lr_cdf
    

def run_iterative_lr(include_hsk_lvl=False, plot_lr=False, num_samples=20,
                     term_absdiff_thres=None, term_max_iter=5, answer_based_on_anki=True, anki_lvl=0):
    # Ask for HSK level
    user_hsk_lvl = anki_lvl
    while user_hsk_lvl == 0:
        print("HSK1: Beginner 1")
        print("HSK2: Beginner 2")
        print("HSK3: Intermediate 1")
        print("HSK4: Intermediate 2")
        print("HSK5: Advanced 1")
        print("HSK6: Advanced 2")

        answer = input('What is roughly your HSK level, 1-6? ')
        fail = False
        try:
            user_hsk_lvl = int(answer)
        except:
            fail = True
        if not (1 <= user_hsk_lvl <= 6):
            fail = True
        if fail:
            print('Please choose a number between 1-6\n')

    sampled_words = set()
    sampled_outcomes = []
    sampled_indices = []
    lr_cdf = None
    stats_history = []
    iteration = 0

    # Set a sampling distribution based on the user's specified HSK level
    # Found that cubing the reverse of the hsk level is a good initial frequency rank coefficient
    # Small levels lead to big coefficients, which lead to a narrower sampling distribution
    # Add a little to the lower levels to make sure low frequency words don't dominate
    initial_sampling_freq_coef = (7 - user_hsk_lvl)**3 + (0.2 if user_hsk_lvl <= 2 else 0.0)
    initial_sampling_intercept = -np.quantile(normalized_data[hsk_lvls == user_hsk_lvl, 0], .75) * initial_sampling_freq_coef

    lr = LogisticRegression()
    lr.coef_ = np.array([initial_sampling_freq_coef, *np.zeros(normalized_data.shape[1]-1)]).reshape(1, -1)
    lr.intercept_ = initial_sampling_intercept
    _, lr_cdf = lr_prediction_and_sampling_cdf(lr)
    lr_coeffs, prev_lr_coeffs = None, None
    def terminate():
        if term_max_iter is not None and iteration >= term_max_iter:
            return True
        if term_absdiff_thres is None or prev_lr_coeffs is None:
            return False
        sum_absdiff = np.abs(lr_coeffs - prev_lr_coeffs).sum()
        return sum_absdiff < term_absdiff_thres

    
    while not terminate():
        if not answer_based_on_anki: print(f'Iteration {iteration+1}')

        curr_sample = []
        while len(curr_sample) < num_samples:
            r = random.random()
            idx = np.searchsorted(lr_cdf, r)
            idx = idx -1 if idx == len(lr_cdf) else idx
            try:
                w = words[idx]
            except:
                breakpoint()
            if (idx in curr_sample or
                    w in sampled_words or
                    w in component_words):
                continue
            curr_sample.append(idx)

        outcomes = []
        sample_words = []
        sample_indices = []
        for idx in curr_sample:
            w, f = words[idx], normalized_data[idx, 0]
            hsk_lvl = hsk_lvl_by_word.get(w, None)

            # Fix race condition in streams resulting in input coming first
            if answer_based_on_anki:
                outcome = 1.0 if w in anki_words else 0.0
            else:
                print(f'HSK {hsk_lvl}: {w} {f}')
                for py, transls in dictionary[w]:
                    print(f'{py}: {"/".join(transls)}')

                time.sleep(0.05)
                answer = input('Do you know it, y/n? ')
                print('') # Give space to the next question
                outcome = 1.0 if answer == 'y' else 0.0

            sampled_words.add(w)
            outcomes.append(outcome)

        sampled_indices = sampled_indices + curr_sample
        sampled_outcomes = sampled_outcomes + outcomes

        if len(np.unique(np.array(sampled_outcomes))) != 2:
            # Need data from both true/false classes before doing regression
            iteration += 1
            continue

        # NOTE: have to set a higher C in order to reduce regularization
        # see https://stackoverflow.com/a/52064154
        # By default C=1.0 which produces very high regularization
        # But we still need some regularization, especially when an HSK1-3 student
        # don't know any words outside HSK, this will make the estimate for rank coefficient
        # too extreme
        lr = LogisticRegression(random_state=0, solver='lbfgs', C=1e2)
        X = normalized_data[np.array(sampled_indices)]
        y = np.array(sampled_outcomes)
        lr.fit(X, y)
        
        prev_lr_coeffs, lr_coeffs = lr_coeffs, np.append(lr.coef_.ravel(), lr.intercept_)
        probs, lr_cdf = lr_prediction_and_sampling_cdf(lr)

        # Evaluate Cross Entropy Loss on words in Anki
        anki_mask = np.array([w in anki_words for w in words])
        anki_loss = -np.log(probs[anki_mask, 1]).sum()
        not_anki_loss = -np.log(probs[~anki_mask, 0]).sum()
        loss = anki_loss + not_anki_loss
        stats = {'anki_loss': anki_loss, 'not_anki_loss': not_anki_loss, 'loss': loss,
                 'outcomes': np.array(outcomes).sum() / len(outcomes), 'coeffs': lr_coeffs}
        for hsk in [True, False]:
            hsk_mask = in_hsk if hsk else ~in_hsk
            pos_outcomes = probs[:, 1] > 0.5
            neg_outcomes = probs[:, 1] < 0.5
            false_negatives = anki_mask & hsk_mask & neg_outcomes
            false_positives = ~anki_mask & hsk_mask & pos_outcomes
            true_negatives = ~anki_mask & hsk_mask & neg_outcomes
            true_positives = anki_mask & hsk_mask & pos_outcomes
            prefix = "hsk_" if hsk else "not_hsk_"
            stats[prefix+'FN'] = np.count_nonzero(false_negatives)
            stats[prefix+'FP'] = np.count_nonzero(false_positives)
            stats[prefix+'TN'] = np.count_nonzero(true_negatives)
            stats[prefix+'TP'] = np.count_nonzero(true_positives)

        stats_history.append(stats)
        iteration += 1
    return stats_history

NUM_ITER = 10
NUM_RUNS = 100

#run_iterative_lr(include_hsk_lvl=True, plot_lr=False, term_max_iter=NUM_ITER,
#                 answer_based_on_anki=False, anki_lvl=2)

loss = []
iteration = []
outcomes = []
final_losses = []
for i in range(NUM_RUNS):
    print('Run ', i)
    stats_history = run_iterative_lr(include_hsk_lvl=True, plot_lr=False, term_max_iter=NUM_ITER,
                                     answer_based_on_anki=True, anki_lvl=5)
    for i, stats in enumerate(stats_history):
        loss.append(stats['loss'])
        outcomes.append(stats['outcomes'])
        iteration.append(i)
    final_losses.append(stats_history[-1]['loss'])
final_losses = np.array(final_losses)
print('Mean final loss:', final_losses.mean(), ' std', final_losses.std())
print('Done')

fig, axes = plt.subplots(1, 2)

sns.lineplot(x=np.array(iteration), y=np.array(loss), ax=axes[0]);
axes[0].set(xlabel='Iteration', ylabel='Loss', title=f'Loss over {NUM_RUNS} runs');
sns.lineplot(x=np.array(iteration), y=np.array(outcomes), ax=axes[1]);
axes[1].set(ylim=(0, 1), xlabel='Iteration', ylabel='% known', title=f'% Known words over {NUM_RUNS} runs');
plt.show()

from IPython.display import HTML, display
import tabulate

h = random.choice(stats_history)
accuracy = []
precision = []
recall = []
for hsk in [True, False]:
    pre = 'hsk_' if hsk else 'not_hsk_'
    FN = h[pre+'FN']
    FP = h[pre+'FP']
    TP = h[pre+'TP']
    TN = h[pre+'TN']
    table = [["HSK" if hsk else "Non-HSK", "Negatives","Positives"],
             ["False", FN, FP],
             ["True", TN, TP]]
    display(HTML(tabulate.tabulate(table, headers="firstrow", tablefmt='html')))
    accuracy.append((TP+TN)/(TP+TN+FP+FN))
    precision.append(float('nan') if TP+FP == 0 else TP/(TP+FP))
    recall.append(TP/(TP+FN))
    
table = [["HSK","Non-HSK"],
         ["Accuracy", accuracy[0], accuracy[1]],
         ["Precision", precision[0], precision[1]],
         ["Recall", recall[0], recall[1]]]
display(HTML(tabulate.tabulate(table, headers="firstrow", tablefmt='html')))