<s> from collections import OrderedDict <EOL> from numpy . linalg import lstsq <EOL> from sklearn . decomposition import LatentDirichletAllocation <EOL> from sklearn . feature_extraction . text import CountVectorizer <EOL> import numpy as np <EOL> import os <EOL> import pandas as pd <EOL> import plac <EOL> def main ( trace_fpath , leaveout = <NUM_LIT> ) : <EOL> leaveout = float ( leaveout ) <EOL> df = pd . read_csv ( trace_fpath , sep = '<STR_LIT:\t>' , names = [ '<STR_LIT>' , '<STR_LIT:u>' , '<STR_LIT:s>' , '<STR_LIT:d>' ] ) <EOL> num_lines = len ( df ) <EOL> to = int ( num_lines - num_lines * leaveout ) <EOL> df_train = df [ : to ] <EOL> df_test = df [ to : ] <EOL> documents_train_right = OrderedDict ( ) <EOL> documents_train_left = OrderedDict ( ) <EOL> tokens_train = set ( ) <EOL> for _ , u , s , d in df_train . values : <EOL> u = str ( u ) <EOL> s = str ( s ) <EOL> d = str ( d ) <EOL> if u not in documents_train_right : <EOL> documents_train_right [ u ] = [ ] <EOL> documents_train_left [ u ] = [ ] <EOL> documents_train_right [ u ] . append ( s ) <EOL> documents_train_left [ u ] . append ( d ) <EOL> tokens_train . add ( s ) <EOL> tokens_train . add ( d ) <EOL> for u in documents_train_right : <EOL> documents_train_right [ u ] = '<STR_LIT:\t>' . join ( documents_train_right [ u ] ) <EOL> documents_train_left [ u ] = '<STR_LIT:\t>' . join ( documents_train_left [ u ] ) <EOL> vectorizer = CountVectorizer ( tokenizer = lambda x : x . split ( '<STR_LIT:\t>' ) , vocabulary = tokens_train ) <EOL> X_train_counts = vectorizer . fit_transform ( documents_train_right . values ( ) ) <EOL> Y_train_counts = vectorizer . transform ( documents_train_left . values ( ) ) <EOL> lda_model = LatentDirichletAllocation ( n_topics = <NUM_LIT:10> , n_jobs = - <NUM_LIT:1> ) <EOL> lda_model . fit ( X_train_counts ) <EOL> Theta_zh = lda_model . transform ( X_train_counts ) . T <EOL> ph = X_train_counts . sum ( axis = <NUM_LIT:1> ) <EOL> pz = np . asarray ( Theta_zh . dot ( ph ) ) [ : , <NUM_LIT:0> ] <EOL> Psi_oz = lda_model . components_ . T <EOL> pz = pz / pz . sum ( ) <EOL> Psi_zo = ( Psi_oz * pz ) . T <EOL> Psi_oz = Psi_oz / Psi_oz . sum ( axis = <NUM_LIT:0> ) <EOL> Psi_zo = Psi_zo / Psi_zo . sum ( axis = <NUM_LIT:0> ) <EOL> X_train_probs = [ ] <EOL> Y_train_probs = [ ] <EOL> for _ , u , s , d in df_train . values : <EOL> if str ( s ) in vectorizer . vocabulary_ and str ( d ) in vectorizer . vocabulary_ : <EOL> id_s = vectorizer . vocabulary_ . get ( str ( s ) ) <EOL> id_d = vectorizer . vocabulary_ . get ( str ( d ) ) <EOL> X_train_probs . append ( Psi_zo [ : , id_s ] ) <EOL> Y_train_probs . append ( Psi_zo [ : , id_d ] ) <EOL> X_train_probs = np . array ( X_train_probs ) <EOL> Y_train_probs = np . array ( Y_train_probs ) <EOL> P_zz = lstsq ( X_train_probs , Y_train_probs ) [ <NUM_LIT:0> ] . T <EOL> P_zz [ P_zz < <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> I = Psi_oz . dot ( P_zz ) <EOL> I = I / I . sum ( axis = <NUM_LIT:0> ) <EOL> probs_tmlda = { } <EOL> probs_lda = { } <EOL> ll_tmlda = <NUM_LIT:0.0> <EOL> ll_lda = <NUM_LIT:0.0> <EOL> n = <NUM_LIT:0> <EOL> for _ , u , s , d in df_test . values : <EOL> u = str ( u ) <EOL> s = str ( s ) <EOL> d = str ( d ) <EOL> if s in vectorizer . vocabulary_ and d in vectorizer . vocabulary_ : <EOL> id_s = vectorizer . vocabulary_ . get ( s ) <EOL> id_d = vectorizer . vocabulary_ . get ( d ) <EOL> if ( id_d , id_s ) not in probs_tmlda : <EOL> probs_tmlda [ id_d , id_s ] = ( Psi_zo [ : , id_s ] * I [ id_s ] ) . sum ( ) <EOL> probs_lda [ id_d , id_s ] = ( Psi_zo [ : , id_s ] * Psi_oz [ id_s ] ) . sum ( ) <EOL> if probs_tmlda [ id_d , id_s ] != <NUM_LIT:0> : <EOL> ll_tmlda += np . log ( probs_tmlda [ id_d , id_s ] ) <EOL> if probs_lda [ id_d , id_s ] != <NUM_LIT:0> : <EOL> ll_lda += np . log ( probs_lda [ id_d , id_s ] ) <EOL> n += <NUM_LIT:1> <EOL> print ( ll_tmlda , ll_lda ) <EOL> print ( ll_tmlda / n , ll_lda / n ) <EOL> print ( n ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> plac . call ( main ) </s>
