1#!/usr/bin/python 2import xgboost as xgb 3from sklearn.datasets import load_svmlight_file 4 5# This script demonstrate how to do ranking with XGBRanker 6x_train, y_train = load_svmlight_file("mq2008.train") 7x_valid, y_valid = load_svmlight_file("mq2008.vali") 8x_test, y_test = load_svmlight_file("mq2008.test") 9 10group_train = [] 11with open("mq2008.train.group", "r") as f: 12 data = f.readlines() 13 for line in data: 14 group_train.append(int(line.split("\n")[0])) 15 16group_valid = [] 17with open("mq2008.vali.group", "r") as f: 18 data = f.readlines() 19 for line in data: 20 group_valid.append(int(line.split("\n")[0])) 21 22group_test = [] 23with open("mq2008.test.group", "r") as f: 24 data = f.readlines() 25 for line in data: 26 group_test.append(int(line.split("\n")[0])) 27 28params = {'objective': 'rank:ndcg', 'learning_rate': 0.1, 29 'gamma': 1.0, 'min_child_weight': 0.1, 30 'max_depth': 6, 'n_estimators': 4} 31model = xgb.sklearn.XGBRanker(**params) 32model.fit(x_train, y_train, group_train, verbose=True, 33 eval_set=[(x_valid, y_valid)], eval_group=[group_valid]) 34pred = model.predict(x_test) 35