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