mq2007.py 10.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
# Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
MQ2007 dataset

MQ2007 is a query set from Million Query track of TREC 2007. There are about 1700 queries in it with labeled documents. In MQ2007, the 5-fold cross
validation strategy is adopted and the 5-fold partitions are included in the package. In each fold, there are three subsets for learning: training set,
D
dzhwinter 已提交
19
validation set and testing set.
20

D
dzhwinter 已提交
21
MQ2007 dataset from website
22 23 24 25 26 27 28
http://research.microsoft.com/en-us/um/beijing/projects/letor/LETOR4.0/Data/MQ2007.rar and parse training set and test set into paddle reader creators

"""

import os
import functools
import rarfile
29
from .common import download
30 31 32 33 34 35 36 37
import numpy as np

# URL = "http://research.microsoft.com/en-us/um/beijing/projects/letor/LETOR4.0/Data/MQ2007.rar"
URL = "http://www.bigdatalab.ac.cn/benchmark/upload/download_source/7b6dbbe2-842c-11e4-a536-bcaec51b9163_MQ2007.rar"
MD5 = "7be1640ae95c6408dab0ae7207bdc706"


def __initialize_meta_info__():
D
dzhwinter 已提交
38
    """
39 40
  download and extract the MQ2007 dataset
  """
D
dzhwinter 已提交
41 42 43 44 45
    fn = fetch()
    rar = rarfile.RarFile(fn)
    dirpath = os.path.dirname(fn)
    rar.extractall(path=dirpath)
    return dirpath
46 47 48


class Query(object):
D
dzhwinter 已提交
49
    """
50 51 52 53 54 55
  queries used for learning to rank algorithms. It is created from relevance scores,  query-document feature vectors

  Parameters:
  ----------
  query_id : int
    query_id in dataset, mapping from query to relevance documents
56
  relevance_score : int
57 58 59 60 61 62 63
    relevance score of query and document pair
  feature_vector : array, dense feature
    feature in vector format
  description : string
    comment section in query doc pair data
  """

D
dzhwinter 已提交
64 65 66 67 68 69 70 71 72 73 74 75
    def __init__(self,
                 query_id=-1,
                 relevance_score=-1,
                 feature_vector=None,
                 description=""):
        self.query_id = query_id
        self.relevance_score = relevance_score
        if feature_vector is None:
            self.feature_vector = []
        else:
            self.feature_vector = feature_vector
        self.description = description
76

D
dzhwinter 已提交
77 78 79 80 81 82 83 84
    def __str__(self):
        string = "%s %s %s" % (str(self.relevance_score), str(self.query_id),
                               " ".join(str(f) for f in self.feature_vector))
        return string

    # @classmethod
    def _parse_(self, text):
        """
85 86
    parse line into Query
    """
D
dzhwinter 已提交
87 88 89 90
        comment_position = text.find('#')
        line = text[:comment_position].strip()
        self.description = text[comment_position + 1:].strip()
        parts = line.split()
91 92 93 94
        if len(parts) != 48:
            sys.stdout.write("expect 48 space split parts, get %d" %
                             (len(parts)))
            return None
95
        # format : 0 qid:10 1:0.000272 2:0.000000 ....
D
dzhwinter 已提交
96 97 98 99 100 101 102
        self.relevance_score = int(parts[0])
        self.query_id = int(parts[1].split(':')[1])
        for p in parts[2:]:
            pair = p.split(':')
            self.feature_vector.append(float(pair[1]))
        return self

103 104

class QueryList(object):
D
dzhwinter 已提交
105
    """
106 107
  group query into list, every item in list is a Query
  """
D
dzhwinter 已提交
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128

    def __init__(self, querylist=None):
        self.query_id = -1
        if querylist is None:
            self.querylist = []
        else:
            self.querylist = querylist
            for query in self.querylist:
                if self.query_id == -1:
                    self.query_id = query.query_id
                else:
                    if self.query_id != query.query_id:
                        raise ValueError("query in list must be same query_id")

    def __iter__(self):
        for query in self.querylist:
            yield query

    def __len__(self):
        return len(self.querylist)

129 130 131
    def __getitem__(self, i):
        return self.querylist[i]

D
dzhwinter 已提交
132 133 134 135 136 137
    def _correct_ranking_(self):
        if self.querylist is None:
            return
        self.querylist.sort(key=lambda x: x.relevance_score, reverse=True)

    def _add_query(self, query):
138
        if self.query_id == -1:
D
dzhwinter 已提交
139
            self.query_id = query.query_id
140
        else:
D
dzhwinter 已提交
141 142 143
            if self.query_id != query.query_id:
                raise ValueError("query in list must be same query_id")
        self.querylist.append(query)
144 145


146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185
def gen_plain_txt(querylist):
    """
  gen plain text in list for other usage
  Paramters:
  --------
  querylist : querylist, one query match many docment pairs in list, see QueryList

  return :
  ------
  query_id : np.array, shape=(samples_num, )
  label : np.array, shape=(samples_num, )
  querylist : np.array, shape=(samples_num, feature_dimension)
    """
    if not isinstance(querylist, QueryList):
        querylist = QueryList(querylist)
    querylist._correct_ranking_()
    for query in querylist:
        yield querylist.query_id, query.relevance_score, np.array(
            query.feature_vector)


def gen_point(querylist):
    """
  gen item in list for point-wise learning to rank algorithm
  Paramters:
  --------
  querylist : querylist, one query match many docment pairs in list, see QueryList

  return :
  ------
  label : np.array, shape=(samples_num, )
  querylist : np.array, shape=(samples_num, feature_dimension)
  """
    if not isinstance(querylist, QueryList):
        querylist = QueryList(querylist)
    querylist._correct_ranking_()
    for query in querylist:
        yield query.relevance_score, np.array(query.feature_vector)


186
def gen_pair(querylist, partial_order="full"):
D
dzhwinter 已提交
187
    """
188 189 190 191 192
  gen pair for pair-wise learning to rank algorithm
  Paramters:
  --------
  querylist : querylist, one query match many docment pairs in list, see QueryList
  pairtial_order : "full" or "neighbour"
193
    there is redudant in all possiable pair combinations, which can be simplifed
194 195 196 197 198 199 200 201
  gen pairs for neighbour items or the full partial order pairs

  return :
  ------
  label : np.array, shape=(1)
  query_left : np.array, shape=(1, feature_dimension)
  query_right : same as left
  """
D
dzhwinter 已提交
202 203 204
    if not isinstance(querylist, QueryList):
        querylist = QueryList(querylist)
    querylist._correct_ranking_()
205 206 207
    labels = []
    docpairs = []

D
dzhwinter 已提交
208
    # C(n,2)
209 210 211 212
    for i in range(len(querylist)):
        query_left = querylist[i]
        for j in range(i + 1, len(querylist)):
            query_right = querylist[j]
D
dzhwinter 已提交
213
            if query_left.relevance_score > query_right.relevance_score:
D
dongzhihong 已提交
214
                labels.append([1])
215 216 217 218 219
                docpairs.append([
                    np.array(query_left.feature_vector),
                    np.array(query_right.feature_vector)
                ])
            elif query_left.relevance_score < query_right.relevance_score:
D
dongzhihong 已提交
220
                labels.append([1])
221 222 223 224 225
                docpairs.append([
                    np.array(query_right.feature_vector),
                    np.array(query_left.feature_vector)
                ])
    for label, pair in zip(labels, docpairs):
D
dongzhihong 已提交
226
        yield np.array(label), pair[0], pair[1]
D
dzhwinter 已提交
227

228 229

def gen_list(querylist):
D
dzhwinter 已提交
230
    """
D
dzhwinter 已提交
231
  gen item in list for list-wise learning to rank algorithm
232 233 234 235 236 237 238 239 240
  Paramters:
  --------
  querylist : querylist, one query match many docment pairs in list, see QueryList

  return :
  ------
  label : np.array, shape=(samples_num, )
  querylist : np.array, shape=(samples_num, feature_dimension)
  """
D
dzhwinter 已提交
241 242
    if not isinstance(querylist, QueryList):
        querylist = QueryList(querylist)
243
    querylist._correct_ranking_()
D
dongzhihong 已提交
244
    relevance_score_list = [[query.relevance_score] for query in querylist]
D
dzhwinter 已提交
245
    feature_vector_list = [query.feature_vector for query in querylist]
D
dongzhihong 已提交
246
    yield np.array(relevance_score_list), np.array(feature_vector_list)
247 248


249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266
def query_filter(querylists):
    """
    filter query get only document with label 0.
    label 0, 1, 2 means the relevance score document with query
    parameters :
      querylist : QueyList list

    return :
      querylist : QueyList list
    """
    filter_query = []
    for querylist in querylists:
        relevance_score_list = [query.relevance_score for query in querylist]
        if sum(relevance_score_list) != .0:
            filter_query.append(querylist)
    return filter_query


D
dzhwinter 已提交
267
def load_from_text(filepath, shuffle=False, fill_missing=-1):
D
dzhwinter 已提交
268
    """
269 270
  parse data file into querys
  """
D
dzhwinter 已提交
271 272 273 274 275 276 277 278
    prev_query_id = -1
    querylists = []
    querylist = None
    fn = __initialize_meta_info__()
    with open(os.path.join(fn, filepath)) as f:
        for line in f:
            query = Query()
            query = query._parse_(line)
279 280
            if query == None:
                continue
D
dzhwinter 已提交
281 282 283 284 285 286
            if query.query_id != prev_query_id:
                if querylist is not None:
                    querylists.append(querylist)
                querylist = QueryList()
                prev_query_id = query.query_id
            querylist._add_query(query)
287 288
    if querylist is not None:
        querylists.append(querylist)
D
dzhwinter 已提交
289
    return querylists
290 291


D
dzhwinter 已提交
292
def __reader__(filepath, format="pairwise", shuffle=False, fill_missing=-1):
D
dzhwinter 已提交
293
    """
294 295 296 297
  Parameters
  --------
  filename : string
  fill_missing : fill the missing value. default in MQ2007 is -1
298

299 300 301 302 303 304
  Returns
  ------
  yield
    label query_left, query_right  # format = "pairwise"
    label querylist # format = "listwise"
  """
305 306 307
    querylists = query_filter(
        load_from_text(
            filepath, shuffle=shuffle, fill_missing=fill_missing))
D
dzhwinter 已提交
308
    for querylist in querylists:
309 310 311 312 313
        if format == "plain_txt":
            yield next(gen_plain_txt(querylist))
        elif format == "pointwise":
            yield next(gen_point(querylist))
        elif format == "pairwise":
D
dzhwinter 已提交
314 315 316 317 318 319 320
            for pair in gen_pair(querylist):
                yield pair
        elif format == "listwise":
            yield next(gen_list(querylist))


train = functools.partial(__reader__, filepath="MQ2007/MQ2007/Fold1/train.txt")
321
test = functools.partial(__reader__, filepath="MQ2007/MQ2007/Fold1/test.txt")
D
dzhwinter 已提交
322

323 324

def fetch():
D
dzhwinter 已提交
325
    return download(URL, "MQ2007", MD5)
326 327


D
dzhwinter 已提交
328 329
if __name__ == "__main__":
    fetch()
330 331 332
    mytest = functools.partial(
        __reader__, filepath="MQ2007/MQ2007/Fold1/sample", format="listwise")
    for label, query in mytest():
333
        print(label, query)