11import unittest
22
33import lightgbm as lgb
4+ import pandas as pd
45
56from common import gpu_test
67
78class TestLightgbm (unittest .TestCase ):
89 # Based on the "simple_example" from their documentation:
910 # https://github.com/Microsoft/LightGBM/blob/master/examples/python-guide/simple_example.py
1011 def test_cpu (self ):
11- lgb_train = lgb .Dataset ('/input/tests/data/lgb_train.bin' )
12- lgb_eval = lgb .Dataset ('/input/tests/data/lgb_test.bin' , reference = lgb_train )
12+ lgb_train , lgb_eval = self .load_datasets ()
1313
1414 params = {
1515 'task' : 'train' ,
@@ -35,9 +35,8 @@ def test_cpu(self):
3535
3636 @gpu_test
3737 def test_gpu (self ):
38- lgb_train = lgb .Dataset ('/input/tests/data/lgb_train.bin' )
39- lgb_eval = lgb .Dataset ('/input/tests/data/lgb_test.bin' , reference = lgb_train )
40-
38+ lgb_train , lgb_eval = self .load_datasets ()
39+
4140 params = {
4241 'boosting_type' : 'gbdt' ,
4342 'objective' : 'regression' ,
@@ -59,3 +58,17 @@ def test_gpu(self):
5958 early_stopping_rounds = 1 )
6059
6160 self .assertEqual (1 , gbm .best_iteration )
61+
62+ def load_datasets (self ):
63+ df_train = pd .read_csv ('/input/tests/data/lgb_train.csv' , header = None , sep = '\t ' )
64+ df_test = pd .read_csv ('/input/tests/data/lgb_test.csv' , header = None , sep = '\t ' )
65+
66+ y_train = df_train [0 ]
67+ y_test = df_test [0 ]
68+ X_train = df_train .drop (0 , axis = 1 )
69+ X_test = df_test .drop (0 , axis = 1 )
70+
71+ lgb_train = lgb .Dataset (X_train , y_train )
72+ lgb_eval = lgb .Dataset (X_test , y_test , reference = lgb_train )
73+
74+ return (lgb_train , lgb_eval )
0 commit comments