<s> import os <EOL> import os . path <EOL> import pickle <EOL> import time <EOL> import chainer <EOL> from chainer import cuda <EOL> from chainer import serializers <EOL> import chainer . optimizers as O <EOL> import numpy as np <EOL> from lda2vec import utils <EOL> from lda2vec import prepare_topics , print_top_words_per_topic <EOL> from lda2vec_model import LDA2Vec <EOL> gpu_id = int ( os . getenv ( '<STR_LIT>' , <NUM_LIT:0> ) ) <EOL> cuda . get_device ( gpu_id ) . use ( ) <EOL> print "<STR_LIT>" + str ( gpu_id ) <EOL> data_dir = os . getenv ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> fn_vocab = '<STR_LIT>' . format ( data_dir = data_dir ) <EOL> fn_corpus = '<STR_LIT>' . format ( data_dir = data_dir ) <EOL> fn_flatnd = '<STR_LIT>' . format ( data_dir = data_dir ) <EOL> fn_docids = '<STR_LIT>' . format ( data_dir = data_dir ) <EOL> vocab = pickle . load ( open ( fn_vocab , '<STR_LIT:r>' ) ) <EOL> corpus = pickle . load ( open ( fn_corpus , '<STR_LIT:r>' ) ) <EOL> flattened = np . load ( fn_flatnd ) <EOL> doc_ids = np . load ( fn_docids ) <EOL> n_docs = doc_ids . max ( ) + <NUM_LIT:1> <EOL> n_vocab = flattened . max ( ) + <NUM_LIT:1> <EOL> clambda = <NUM_LIT> <EOL> n_topics = int ( os . getenv ( '<STR_LIT>' , <NUM_LIT:20> ) ) <EOL> batchsize = <NUM_LIT> <EOL> power = float ( os . getenv ( '<STR_LIT>' , <NUM_LIT> ) ) <EOL> pretrained = bool ( os . getenv ( '<STR_LIT>' , True ) ) <EOL> n_units = int ( os . getenv ( '<STR_LIT>' , <NUM_LIT> ) ) <EOL> words = corpus . word_list ( vocab ) [ : n_vocab ] <EOL> doc_idx , lengths = np . unique ( doc_ids , return_counts = True ) <EOL> doc_lengths = np . zeros ( doc_ids . max ( ) + <NUM_LIT:1> , dtype = '<STR_LIT>' ) <EOL> doc_lengths [ doc_idx ] = lengths <EOL> tok_idx , freq = np . unique ( flattened , return_counts = True ) <EOL> term_frequency = np . zeros ( n_vocab , dtype = '<STR_LIT>' ) <EOL> term_frequency [ tok_idx ] = freq <EOL> model = LDA2Vec ( n_documents = n_docs , n_document_topics = n_topics , <EOL> n_units = n_units , n_vocab = n_vocab , counts = term_frequency , <EOL> n_samples = <NUM_LIT:15> , power = power ) <EOL> if os . path . exists ( '<STR_LIT>' ) : <EOL> print "<STR_LIT>" <EOL> serializers . load_hdf5 ( "<STR_LIT>" , model ) <EOL> model . to_gpu ( ) <EOL> optimizer = O . Adam ( ) <EOL> optimizer . setup ( model ) <EOL> clip = chainer . optimizer . GradientClipping ( <NUM_LIT> ) <EOL> optimizer . add_hook ( clip ) <EOL> j = <NUM_LIT:0> <EOL> epoch = <NUM_LIT:0> <EOL> fraction = batchsize * <NUM_LIT:1.0> / flattened . shape [ <NUM_LIT:0> ] <EOL> for epoch in range ( <NUM_LIT> ) : <EOL> data = prepare_topics ( cuda . to_cpu ( model . mixture . weights . W . data ) . copy ( ) , <EOL> cuda . to_cpu ( model . mixture . factors . W . data ) . copy ( ) , <EOL> cuda . to_cpu ( model . sampler . W . data ) . copy ( ) , <EOL> words ) <EOL> print_top_words_per_topic ( data ) <EOL> data [ '<STR_LIT>' ] = doc_lengths <EOL> data [ '<STR_LIT>' ] = term_frequency <EOL> np . savez ( '<STR_LIT>' , ** data ) <EOL> for d , f in utils . chunks ( batchsize , doc_ids , flattened ) : <EOL> t0 = time . time ( ) <EOL> optimizer . zero_grads ( ) <EOL> l = model . fit_partial ( d . copy ( ) , f . copy ( ) ) <EOL> prior = model . prior ( ) <EOL> loss = prior * fraction <EOL> loss . backward ( ) <EOL> optimizer . update ( ) <EOL> msg = ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> prior . to_cpu ( ) <EOL> loss . to_cpu ( ) <EOL> t1 = time . time ( ) <EOL> dt = t1 - t0 <EOL> rate = batchsize / dt <EOL> logs = dict ( loss = float ( l ) , epoch = epoch , j = j , <EOL> prior = float ( prior . data ) , rate = rate ) <EOL> print msg . format ( ** logs ) <EOL> j += <NUM_LIT:1> <EOL> serializers . save_hdf5 ( "<STR_LIT>" , model ) </s>
