<s> import ffm <EOL> import numpy as np <EOL> from sklearn . metrics import mean_squared_error <EOL> from sklearn . utils import assert_all_finite <EOL> from . validation import ( assert_all_finite , check_consistent_length , <EOL> check_array ) <EOL> from . base import ( FactorizationMachine , _validate_class_labels , <EOL> _check_warm_start ) <EOL> def find_init_stdev ( fm , X_train , y_train , X_vali = None , y_vali = None , <EOL> stdev_range = None , ) : <EOL> if not stdev_range : <EOL> stdev_range = [ <NUM_LIT:0.1> , <NUM_LIT:0.1> , <NUM_LIT> , <NUM_LIT:0.5> , <NUM_LIT:1.0> ] <EOL> if not isinstance ( fm , FMRegression ) : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> if X_vali is None : <EOL> X_test = X_train [ : <NUM_LIT:2> , : ] <EOL> else : <EOL> X_test = X_vali <EOL> best_init_stdev = <NUM_LIT:0> <EOL> best_mse = np . finfo ( np . float64 ) . max <EOL> for init_stdev in stdev_range : <EOL> fm . init_stdev = init_stdev <EOL> y_pred_vali = fm . fit_predict ( X_train , y_train , X_test ) <EOL> if X_vali is None : <EOL> y_pred = fm . predict ( X_train ) <EOL> mse = mean_squared_error ( y_pred , y_train ) <EOL> else : <EOL> mse = mean_squared_error ( y_pred_vali , y_vali ) <EOL> if mse < best_mse : <EOL> best_mse = mse <EOL> best_init_stdev = init_stdev <EOL> return best_init_stdev , best_mse <EOL> def _validate_mcmc_fit_input ( X_train , y_train , X_test ) : <EOL> check_consistent_length ( X_train , y_train ) <EOL> assert_all_finite ( y_train ) <EOL> y_train = check_array ( y_train , ensure_2d = False , dtype = np . float64 ) <EOL> assert X_train . shape [ <NUM_LIT:1> ] == X_test . shape [ <NUM_LIT:1> ] <EOL> X_train = check_array ( X_train , accept_sparse = "<STR_LIT>" , dtype = np . float64 , <EOL> order = "<STR_LIT:F>" ) <EOL> X_test = check_array ( X_test , accept_sparse = "<STR_LIT>" , dtype = np . float64 , <EOL> order = "<STR_LIT:F>" ) <EOL> return X_train , y_train , X_test <EOL> class FMRegression ( FactorizationMachine ) : <EOL> """<STR_LIT>""" <EOL> def fit_predict ( self , X_train , y_train , X_test , n_more_iter = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> self . task = "<STR_LIT>" <EOL> X_train , y_train , X_test = _validate_mcmc_fit_input ( X_train , y_train , <EOL> X_test ) <EOL> self . n_iter = self . n_iter + n_more_iter <EOL> if n_more_iter > <NUM_LIT:0> : <EOL> _check_warm_start ( self , X_train ) <EOL> assert self . prediction_ . shape [ <NUM_LIT:0> ] == X_test . shape [ <NUM_LIT:0> ] <EOL> assert self . hyper_param_ . shape <EOL> self . warm_start = True <EOL> else : <EOL> self . iter_count = <NUM_LIT:0> <EOL> coef , y_pred = ffm . ffm_mcmc_fit_predict ( self , X_train , <EOL> X_test , y_train ) <EOL> self . w0_ , self . w_ , self . V_ = coef <EOL> self . prediction_ = y_pred <EOL> self . warm_start = False <EOL> if self . iter_count != <NUM_LIT:0> : <EOL> self . iter_count = self . iter_count + n_more_iter <EOL> else : <EOL> self . iter_count = self . n_iter <EOL> return y_pred <EOL> class FMClassification ( FactorizationMachine ) : <EOL> """<STR_LIT>""" <EOL> def fit_predict ( self , X_train , y_train , X_test ) : <EOL> """<STR_LIT>""" <EOL> y_proba = self . fit_predict_proba ( X_train , y_train , X_test ) <EOL> y_pred = np . zeros_like ( y_proba , dtype = np . float64 ) + self . classes_ [ <NUM_LIT:0> ] <EOL> y_pred [ y_proba > <NUM_LIT> ] = self . classes_ [ <NUM_LIT:1> ] <EOL> return y_pred <EOL> def fit_predict_proba ( self , X_train , y_train , X_test ) : <EOL> """<STR_LIT>""" <EOL> self . task = "<STR_LIT>" <EOL> self . classes_ = np . unique ( y_train ) <EOL> if len ( self . classes_ ) != <NUM_LIT:2> : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % self . classes_ ) <EOL> y_train = y_train . copy ( ) <EOL> i_class1 = ( y_train == self . classes_ [ <NUM_LIT:0> ] ) <EOL> y_train [ i_class1 ] = - <NUM_LIT:1> <EOL> y_train [ - i_class1 ] = <NUM_LIT:1> <EOL> X_train , y_train , X_test = _validate_mcmc_fit_input ( X_train , y_train , <EOL> X_test ) <EOL> y_train = _validate_class_labels ( y_train ) <EOL> coef , y_pred = ffm . ffm_mcmc_fit_predict ( self , X_train , <EOL> X_test , y_train ) <EOL> self . w0_ , self . w_ , self . V_ = coef <EOL> return y_pred </s>
