import numpy as np
import xgboost as xgb
from xgboost import XGBRegressor, DMatrix
import warnings
warnings.filterwarnings('ignore')

text_columns = ['query_id', 'region', 'relevance', 'org_id', 'window_center', 'query', 'query_tokens', 'all_text_info', 'org_name']

X_train = train_df.drop(text_columns, axis=1)
y_train = train_df['relevance']
queries_train = train_df['query_id'].values
xgtrain = DMatrix(X_train, label=y_train)
xgtrain.set_group(train_df.groupby('relevance').apply(len).values)

X_test = test_df.drop(text_columns, axis=1)
y_test = test_df['relevance'].values
queries_test = test_df['query_id'].values
xgtest = DMatrix(X_test, label=y_test)
xgtest.set_group(test_df.groupby('relevance').apply(len).values)

bst = xgb.train({
    "objective": "rank:pairwise",
    "eval_metric": "ndcg@10",
    "silent": True
}, xgtrain, num_boost_round=50)

print(bst.eval(xgtest))