MSMARCO support: monoBERT (#14)

* add monobert for marco

* temp change python3.7 to 3.6 for colab compatibility

* fix evaluation options

* fix issues

* add missing options in evaluate_passage_ranker

* working monobert

* update transformers, clean code

* update tokenizers

* add dataclasses if < 3.7

* cleanup todos

* update to newer transformers along with syntax, clean up settings

* model-name-or-path as str type

* fix tokenizer loading for t5

pygaggle/data/__init__.py

@@ -1,2 +1,3 @@
 from .kaggle import *
 from .relevance import *
+from .msmarco import *

pygaggle/data/msmarco.py

@@ -0,0 +1,137 @@
+import os
+from collections import OrderedDict, defaultdict
+from typing import List, Set, DefaultDict
+import json
+import logging
+from itertools import permutations
+
+from pydantic import BaseModel
+import scipy.special as sp
+import numpy as np
+
+from .relevance import RelevanceExample, MsMarcoPassageLoader
+from pygaggle.model.tokenize import SpacySenticizer
+from pygaggle.rerank.base import Query, Text
+from pygaggle.data.unicode import convert_to_unicode
+
+
+__all__ = ['MsMarcoExample', 'MsMarcoDataset']
+
+
+class MsMarcoExample(BaseModel):
+ qid: str
+ text: str
+ candidates: List[str]
+ relevant_candidates: Set[str]
+
+class MsMarcoDataset(BaseModel):
+ examples: List[MsMarcoExample]
+
+ @classmethod
+ def load_qrels(cls, path: str) -> DefaultDict[str, Set[str]]:
+ qrels = defaultdict(set)
+ with open(path) as f:
+ for i, line in enumerate(f):
+ qid, _, doc_id, relevance = line.rstrip().split('\t')
+ if int(relevance) >= 1:
+ qrels[qid].add(doc_id)
+ return qrels
+
+ @classmethod
+ def load_run(cls, path: str):
+ '''Returns OrderedDict[str, List[str]]'''
+ run = OrderedDict()
+ with open(path) as f:
+ for i, line in enumerate(f):
+ qid, doc_title, rank = line.split('\t')
+ if qid not in run:
+ run[qid] = []
+ run[qid].append((doc_title, int(rank)))
+ sorted_run = OrderedDict()
+ for qid, doc_titles_ranks in run.items():
+ sorted(doc_titles_ranks, key=lambda x: x[1])
+ doc_titles = [doc_titles for doc_titles, _ in doc_titles_ranks]
+ sorted_run[qid] = doc_titles
+ return sorted_run
+
+ @classmethod
+ def load_queries(cls, 
+ path: str, 
+ qrels: DefaultDict[str, Set[str]], 
+ run) -> List[MsMarcoExample]:
+ queries = []
+ with open(path) as f:
+ for i, line in enumerate(f):
+ qid, query = line.rstrip().split('\t')
+ candidates = run[qid]
+ queries.append(MsMarcoExample(qid = qid,
+ text = query,
+ candidates = run[qid],
+ relevant_candidates = qrels[qid]))
+ return queries
+
+ @classmethod
+ def from_folder(cls, 
+ folder: str, 
+ split: str = 'dev', 
+ is_duo: bool = False) -> 'MsMarcoDataset':
+ run_mono = "mono." if is_duo else ""
+ query_path = os.path.join(folder, f"queries.{split}.small.tsv")
+ qrels_path = os.path.join(folder, f"qrels.{split}.small.tsv")
+ run_path = os.path.join(folder, f"run.{run_mono}{split}.small.tsv")
+ return cls(examples = cls.load_queries(query_path, 
+ cls.load_qrels(qrels_path),
+ cls.load_run(run_path)))
+
+
+ def query_passage_tuples(self, is_duo: bool = False):
+ return (((ex.qid, ex.text, ex.relevant_candidates), perm_pas) for ex in self.examples
+ for perm_pas in permutations(ex.candidates, r=1+int(is_duo)))
+
+
+ def to_relevance_examples(self,
+ index_path: str,
+ is_duo: bool = False) -> List[RelevanceExample]:
+ loader = MsMarcoPassageLoader(index_path)
+ example_map = {}
+ for (qid, text, rel_cands), cands in self.query_passage_tuples():
+ if qid not in example_map:
+ example_map[qid] = [convert_to_unicode(text), [], [], []]
+ example_map[qid][1].append([cand for cand in cands][0])
+ try:
+ passages = [loader.load_passage(cand) for cand in cands]
+ example_map[qid][2].append([convert_to_unicode(passage.all_text) for passage in passages][0]) 
+ except ValueError as e:
+ logging.warning(f'Skipping {passages}')
+ continue
+ example_map[qid][3].append(cands[0] in rel_cands)
+ mean_stats = defaultdict(list)
+ for ex in self.examples:
+ int_rels = np.array(list(map(int, example_map[ex.qid][3])))
+ p = int_rels.sum()/(len(ex.candidates) - 1) if is_duo else int_rels.sum()
+ mean_stats['Random P@1'].append(np.mean(int_rels))
+ n = len(ex.candidates) - p
+ N = len(ex.candidates)
+ if len(ex.candidates) <= 1000:
+ mean_stats['Random R@1000'].append(1 if 1 in int_rels else 0)
+ numer = np.array([sp.comb(n, i) / (N - i) for i in range(0, n + 1) if i!=N]) * p
+ if n == N:
+ numer = np.append(numer, 0)
+ denom = np.array([sp.comb(N, i) for i in range(0, n + 1)])
+ rr = 1 / np.arange(1, n + 2)
+ rmrr = np.sum(numer * rr / denom)
+ mean_stats['Random MRR'].append(rmrr)
+ rmrr10 = np.sum(numer[:10] * rr[:10] / denom[:10])
+ mean_stats['Random MRR@10'].append(rmrr10)
+ ex_index = len(ex.candidates)
+ for rel_cand in ex.relevant_candidates:
+ if rel_cand in ex.candidates:
+ ex_index = min(ex.candidates.index(rel_cand), ex_index)
+ mean_stats['Existing MRR'].append(1 / (ex_index + 1) if ex_index < len(ex.candidates) else 0)
+ mean_stats['Existing MRR@10'].append(1 / (ex_index + 1) if ex_index < 10 else 0)
+ for k, v in mean_stats.items():
+ logging.info(f'{k}: {np.mean(v)}')
+ return [RelevanceExample(Query(text=query_text, id=qid), 
+ list(map(lambda s: Text(s[1], dict(docid=s[0])), zip(cands, cands_text))), 
+ rel_cands) \
+ for qid, (query_text, cands, cands_text, rel_cands) in example_map.items()]

pygaggle/data/relevance.py

@@ -30,6 +30,15 @@ def all_text(self):
  return '\n'.join((self.abstract, self.body_text, self.ref_entries))
 
 
+@dataclass
+class MsMarcoPassage:
+ para_text: str
+
+ @property
+ def all_text(self):
+ return self.para_text
+
+
 class Cord19DocumentLoader:
  double_space_pattern = re.compile(r'\s\s+')
 
@@ -50,3 +59,15 @@ def unfold(entries):
  return Cord19Document(unfold(article['abstract']),
  unfold(article['body_text']),
  unfold(ref_entries))
+
+
+class MsMarcoPassageLoader:
+ def __init__(self, index_path: str):
+ self.searcher = pysearch.SimpleSearcher(index_path)
+
+ def load_passage(self, id: str) -> MsMarcoPassage:
+ try:
+ passage = self.searcher.doc(id).lucene_document().get('raw')
+ except AttributeError:
+ raise ValueError('passage unretrievable')
+ return MsMarcoPassage(passage)

pygaggle/data/unicode.py

@@ -0,0 +1,8 @@
+def convert_to_unicode(text):
+ """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
+ if isinstance(text, str):
+ return text
+ elif isinstance(text, bytes):
+ return text.decode("utf-8", "ignore")
+ else:
+ raise ValueError("Unsupported string type: %s" % (type(text)))

pygaggle/model/decode.py

@@ -19,7 +19,10 @@ def greedy_decode(model: PreTrainedModel,
  past = model.get_encoder()(input_ids, attention_mask=attention_mask)
  next_token_logits = None
  for _ in range(length):
- model_inputs = model.prepare_inputs_for_generation(decode_ids, past=past, attention_mask=attention_mask)
+ model_inputs = model.prepare_inputs_for_generation(decode_ids,
+ past=past, 
+ attention_mask=attention_mask, 
+ use_cache=True)
  outputs = model(**model_inputs) # (batch_size, cur_len, vocab_size)
  next_token_logits = outputs[0][:, -1, :] # (batch_size, vocab_size)
  decode_ids = torch.cat([decode_ids, next_token_logits.max(1)[1].unsqueeze(-1)], dim=-1)

pygaggle/model/evaluate.py

@@ -100,6 +100,16 @@ class RecallAt3Metric(TopkMixin, RecallAccumulator):
  top_k = 3
 
 
+@register_metric('recall@50')
+class RecallAt50Metric(TopkMixin, RecallAccumulator):
+ top_k = 50
+
+
+@register_metric('recall@1000')
+class RecallAt1000Metric(TopkMixin, RecallAccumulator):
+ top_k = 1000
+
+
 @register_metric('mrr')
 class MrrMetric(MeanAccumulator):
  def accumulate(self, scores: List[float], gold: RelevanceExample):
@@ -108,6 +118,13 @@ def accumulate(self, scores: List[float], gold: RelevanceExample):
  self.scores.append(rr)
 
 
+@register_metric('mrr@10')
+class MrrAt10Metric(MeanAccumulator):
+ def accumulate(self, scores: List[float], gold: RelevanceExample):
+ scores = sorted(list(enumerate(scores)), key=lambda x: x[1], reverse=True)
+ rr = next((1 / (rank_idx + 1) for rank_idx, (idx, _) in enumerate(scores) if (gold.labels[idx] and rank_idx < 10)), 0)
+ self.scores.append(rr)
+
 class ThresholdedRecallMetric(DynamicThresholdingMixin, RecallAccumulator):
  threshold = 0.5
 

pygaggle/rerank/base.py

@@ -17,9 +17,12 @@ class Query:
  ----------
  text : str
  The query text.
+ id : Optional[str]
+ The query id.
  """
- def __init__(self, text: str):
+ def __init__(self, text: str, id: Optional[str] = None):
  self.text = text
+ self.id = id
 
 
 class Text:

pygaggle/run/evaluate_kaggle_highlighter.py

@@ -16,10 +16,10 @@
 from pygaggle.rerank.similarity import CosineSimilarityMatrixProvider
 from pygaggle.model import SimpleBatchTokenizer, CachedT5ModelLoader, T5BatchTokenizer, RerankerEvaluator, metric_names
 from pygaggle.data import LitReviewDataset
-from pygaggle.settings import Settings
+from pygaggle.settings import Cord19Settings
 
 
-SETTINGS = Settings()
+SETTINGS = Cord19Settings()
 METHOD_CHOICES = ('transformer', 'bm25', 't5', 'seq_class_transformer', 'qa_transformer', 'random')