# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
from data import * 
from lda_trial import *

from sklearn.linear_model import SGDClassifier
from sklearn.svm import LinearSVC
from sklearn.metrics import accuracy_score
# -

trainset = IMDBDataset('train', data_limit=20_000)
validset = IMDBDataset('valid')

train_accu

lens = [len(x) for x in trainset.x]

np.array(lens).mean()

# +
n = 20000

ks = [2, 5, 10, 20, 50]

train_accu = [0 for _ in range(len(ks))]
test_accu = [0 for _ in range(len(ks))]
times = []

for _ in range(3):
    for idx, k in enumerate(ks):
        print(f'{k}.')
        start = time.time()
        trained_model, final_metric = tp_one_trial(trainset, 'slda', k, n, 
                                                  3, 5,  # args.burn_in,
                                                  max_iter=1000, stop_increase=5, metric='ll')
        times.append(time.time() - start)
        lda_x, lda_y = load_LDA_data_batch(trained_model, trainset)
        model = LinearSVC()
        model.fit(lda_x, lda_y)

        prediction = model.predict(lda_x)
        ground_truth = lda_y
        print(f'Train: {100*accuracy_score(prediction, ground_truth):6.5f}%')
        train_accu[idx] += accuracy_score(prediction, ground_truth)

        lda_x, lda_y = load_LDA_data_batch(trained_model, validset)
        prediction = model.predict(lda_x)
        ground_truth = lda_y
        print(f'Test: {100*accuracy_score(prediction, ground_truth):6.5f}%')
        test_accu[idx] += accuracy_score(prediction, ground_truth)
        print('-' * 100)
# -

n = 20000
train_accu = []
test_accu = []
for k in [2, 3, 5, 10, 20, 50, 100]:
    print(f'{k}.')
    trained_model, final_metric = tp_one_trial(trainset, 'slda', k, n, 
                                              3, 5,  # args.burn_in,
                                              max_iter=20, stop_increase=5, metric='ll')

    lda_x, lda_y = load_LDA_data_batch(trained_model, trainset)
    model = LinearSVC()
    model.fit(lda_x, lda_y)

    prediction = model.predict(lda_x)
    ground_truth = lda_y
    print(f'Train: {100*accuracy_score(prediction, ground_truth):6.5f}%')
    train_accu.append(accuracy_score(prediction, ground_truth))
    lda_x, lda_y = load_LDA_data_batch(trained_model, validset)

    prediction = model.predict(lda_x)
    ground_truth = lda_y
    print(f'Test: {100*accuracy_score(prediction, ground_truth):6.5f}%')
    train_accu.append(accuracy_score(prediction, ground_truth))
    print('-' * 100)




StarcoderdataJupyterScriptsDedupFiltered