# ---
# 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
# ---

# +
import pandas as pd
from glob import glob
from os import listdir, chdir
from re import compile as rcompile
from path import Path
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer

from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam, Nadam

from sklearn.metrics.pairwise import cosine_similarity

# +
trex = rcompile('[^a-zA-Z 0-9]')
tokenize = lambda x: trex.sub('', x.lower().replace(',', ' ').replace('-', ' '))

def read_transcript(file_name):
    file_names = glob(f'{file_name}*')
    ret = ''
    for fname in file_names:
        with open(fname) as fd:
            ret += fd.read()+'\n\n'
    return ret

reference_transcripts_dir = '../data/transcripts'
google_transcripts_dir = '../data/google_transcripts'

reference_transcripts_files = listdir(reference_transcripts_dir)
google_transcripts_files = listdir(google_transcripts_dir)
reference_transcripts_files.sort()
google_transcripts_files.sort()
# -

print(reference_transcripts_files)
print(google_transcripts_files)

#Some of the reference transcripts have missing information and those files have
#filename with '(missing page <number>)' substrings - filter these files out for
#now.
rex = rcompile('[0-9]+.txt')
reference_transcripts_file_roots = [x.split('.')[0] for x in reference_transcripts_files if rex.match(x) is not None]
print(reference_transcripts_file_roots)

# +

with Path(reference_transcripts_dir):
    reference_transcripts = [read_transcript(f'{x}.txt') for x in reference_transcripts_file_roots]

print(len(reference_transcripts),'\n', reference_transcripts[0],'\n', reference_transcripts[1])

with Path(google_transcripts_dir):
    google_transcripts = [read_transcript(x) for x in reference_transcripts_file_roots]
    
print(len(google_transcripts),'\n', google_transcripts[0],'\n', google_transcripts[1])

reference_transcripts = list(map(tokenize, reference_transcripts))
print(len(reference_transcripts),'\n', reference_transcripts[0],'\n', reference_transcripts[1])

google_transcripts = list(map(tokenize, google_transcripts))
print(len(google_transcripts),'\n', google_transcripts[0],'\n', google_transcripts[1])


# +
# Instantiate vectorizer object
tfidf = TfidfVectorizer(stop_words='english',
                        ngram_range=(1,2),
                        min_df=1
                       )
# Create a vocabulary and get word counts per document
ref_dtm = tfidf.fit_transform(reference_transcripts)

features = tfidf.get_feature_names()
#display(len(features), features[:50])

# Get feature names to use as dataframe column headers
ref_dtm = pd.DataFrame(ref_dtm.todense(), columns=features)
print(ref_dtm.shape)
display(ref_dtm.head())


# +
ishape = ref_dtm.shape[1]
# Create Model 
input_img = Input(shape=(ishape, ))

x = Dense(512)(input_img)

x = Dense(256)(x)

x = Dense(128)(x)

encoded = Dense(64)(x)

x = Dense(128)(encoded)

x = Dense(256)(x)

x = Dense(512, activation='sigmoid')(x)
decoded = Dense(ishape, activation='sigmoid')(x)


rmodel = Model(input_img, decoded)
rmodel.compile(loss='mse', optimizer=Adam(learning_rate=0.01))

rmodel.fit(ref_dtm, ref_dtm, batch_size=512, epochs=5)

# +
encoder = Model(input_img, encoded)

encoded_ref_dtm = encoder.predict(ref_dtm)
print(encoded_ref_dtm.shape)
display(encoded_ref_dtm)


# -

def get_vector(doc):
    vec = tfidf.transform([tokenize(doc)]).todense()
    return encoder.predict(vec)


print(get_vector(google_transcripts[0]))


# +
def get_encoded_dtm(transcripts):
    dtm = tfidf.transform(transcripts).todense()
    print(dtm.shape)
    encoded_dtm = encoder.predict(dtm)
    print(encoded_dtm.shape)
    display(encoded_dtm)
    return encoded_dtm

encoded_google_dtm = get_encoded_dtm(google_transcripts)
print(encoded_google_dtm.shape)
display(encoded_google_dtm)
# -

cosine_similarity([encoded_ref_dtm[0]], [encoded_google_dtm[0]])


# +
def get_cossim(encoded_dtm):
    cossim = []
    for idx,ref in enumerate(encoded_ref_dtm):
        cossim.append(cosine_similarity([ref], [encoded_dtm[idx]])[0][0]) 
    return cossim

cossim = get_cossim(encoded_google_dtm)
print(cossim)
mean_cossim = sum(cossim)/len(cossim)
print(mean_cossim, round((1-mean_cossim)*1e5,4))

# +
from joblib import dump
MODELS_DIR = '../models/'

TFIDF = MODELS_DIR+'tfidf.pkl'
ENCODED_REF_DTM = MODELS_DIR+'encoded_ref_dtm.pkl'
ENCODER = MODELS_DIR+'encoder.h5'
FILE_ROOTS = MODELS_DIR+'file_roots.pkl'

dump(tfidf, TFIDF)
encoder.save(ENCODER)
dump(encoded_ref_dtm, ENCODED_REF_DTM)
dump(reference_transcripts_file_roots, FILE_ROOTS)

StarcoderdataJupyterScriptsDedupFiltered