semaj83/ctmatch

1 2from typing import Dict, List, Tuple3 4from sklearn.metrics import f1_score5from collections import defaultdict6from lxml import etree7import numpy as np8 9  10def get_trec_topic2text(topic_path) -> Dict[str, str]:11    """12    desc:       main method for processing a single XML file of TREC21 patient descriptions called "topics" in this sense13    returns:    dict of topicid: topic text14    """15 16    topic2text = {}17    topic_root = etree.parse(topic_path).getroot()18    for topic in topic_root:19        topic2text[topic.attrib['number']] = topic.text20    21    return topic2text22    23 24 25def get_kz_topic2text(topic_path) -> Dict[str, str]:26    """27    desc:       main method for processing a single XML file of TREC21 patient descriptions called "topics" in this sense28    returns:    dict of topicid: topic text29    """30 31    topic2text = {}32    with open(topic_path, 'r') as f:33        for line in f.readlines():34            line = line.strip()35            36            if line.startswith('<TOP>'):37                topic_id, text = None, None38                continue39 40            if line.startswith('<NUM>'):41                topic_id = line[5:-6]42 43            elif line.startswith('<TITLE>'):44                text = line[7:].strip()45                topic2text[topic_id] = text46 47    return topic2text48          49 50 51def calc_first_positive_rank(ranked_ids: List[str], doc2rel: Dict[str, int], pos_val: int = 2) -> Tuple[int, float]:52    """53    desc:       compute the mean reciprocal rank of a ranking54    returns:    mrr 55    """56    for i, doc_id in enumerate(ranked_ids):57        if doc2rel[doc_id] == pos_val:58            return i + 1, 1./float(i+1)59    return len(ranked_ids) + 1, 0.060 61 62def calc_f1(ranked_ids: List[str], doc2rel: Dict[str, int]) -> Dict[str, Dict[str, float]]:63  label_counts = get_label_counts(doc2rel)64  predicted, ground_truth = [], []65  for doc_id in ranked_ids:66    # 2, 1, 067    ground_truth.append(doc2rel[doc_id])68    pred_label = get_predicted_label(label_counts)69    predicted.append(pred_label)70    label_counts[pred_label] -= 171 72  return f1_score(ground_truth, predicted, average='micro')73 74  75 76def get_label_counts(doc2rel: Dict[str, int]) -> Dict[int, int]:77  """78  return an ordered list of [(2, <count_2s>), (1, <count_1s>), (0, count_0s)]79  """80  label_counts = defaultdict(int)81  for scored_doc in doc2rel:82    label = doc2rel[scored_doc]83    label_counts[label] += 184  return label_counts85 86def get_predicted_label(label_counts: Dict[int, int]) -> int:87  if label_counts[2] > 0:88    return 289  if label_counts[1] > 0:90    return 191  return 0