Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update demo for ranking. (#5154)

Browse Source
This commit is contained in:
Jiaming Yuan
2019-12-24 13:39:07 +08:00
committed by GitHub
parent 0202e04a8e
commit 73b1bd2789
3 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
demo/rank/rank_sklearn.py
View File

@@ -2,7 +2,6 @@
import xgboost as xgb
from sklearn.datasets import load_svmlight_file
# This script demonstrate how to do ranking with XGBRanker
x_train, y_train = load_svmlight_file("mq2008.train")
x_valid, y_valid = load_svmlight_file("mq2008.vali")
@@ -26,10 +25,10 @@ with open("mq2008.test.group", "r") as f:
for line in data:
group_test.append(int(line.split("\n")[0]))
params = {'objective': 'rank:pairwise', 'learning_rate': 0.1,
params = {'objective': 'rank:ndcg', 'learning_rate': 0.1,
'gamma': 1.0, 'min_child_weight': 0.1,
'max_depth': 6, 'n_estimators': 4}
model = xgb.sklearn.XGBRanker(**params)
model.fit(x_train, y_train, group_train,
model.fit(x_train, y_train, group_train, verbose=True,
eval_set=[(x_valid, y_valid)], eval_group=[group_valid])
pred = model.predict(x_test)