diff --git a/medcat-v2/medcat/stats/stats.py b/medcat-v2/medcat/stats/stats.py index 2b6c6a15d..283a3ef40 100644 --- a/medcat-v2/medcat/stats/stats.py +++ b/medcat-v2/medcat/stats/stats.py @@ -1,309 +1,984 @@ -from typing import Optional, Callable, cast +from typing import Optional, Callable, TextIO, TypedDict from tqdm import tqdm -import traceback from medcat.cat import CAT from medcat.utils.filters import project_filters from medcat.data.mctexport import ( MedCATTrainerExport, MedCATTrainerExportProject, - MedCATTrainerExportDocument, MedCATTrainerExportAnnotation) + MedCATTrainerExportDocument) from medcat.config.config import LinkingFilters from medcat.cdb.concepts import CUIInfo, get_new_cui_info -from medcat.tokenizing.tokens import MutableEntity, MutableDocument +from medcat.tokenizing.tokens import MutableEntity, UNTOKENIZABLE_ENTITY_ID +from medcat.components.types import CoreComponentType +from medcat.utils.training_utils import dataset_aware_component +from collections import defaultdict +from pydantic import BaseModel, Field +from enum import Enum +class MetricMode(str, Enum): + """Supported evaluation modes for statistics collection.""" -class StatsBuilder: + FULL = "full_pipeline" + NER = "ner" + LINKING = "linking" + + +class RawStats(BaseModel): + """Raw accumulated state for a single evaluation mode.""" + + tp: int = 0 + fp: int = 0 + fn: int = 0 + # Number of labels where it is not possible to generate an entity + # Because the tokenizer doesn't find a single token + # i.e. entity at chars 100-103, token is 100-104. + no_tokens: int = 0 + + # per document IoU metrics, summed over all documents and + # averaged later via char_docs, which counts the number of + # documents that have entities that have been processed. + iou_sum: float = 0.0 + giou_sum: float = 0.0 + cohen_k_sum: float = 0.0 + char_docs: int = 0 + + # metrics for individual CUIs + cui_tp: dict[str, int] = Field(default_factory=dict) + cui_fp: dict[str, int] = Field(default_factory=dict) + cui_fn: dict[str, int] = Field(default_factory=dict) + # gold counts is the number of labels for that CUI + cui_gold_counts: dict[str, int] = Field(default_factory=dict) + cui_no_tokens: dict[str, int] = Field(default_factory=dict) + + examples: dict[str, dict[str, list]] = { + 'tp': {}, 'fp': {}, 'fn': {}} + + # character metrics for individual CUIs + cui_iou: defaultdict[str, list[float]] = Field( + default_factory=lambda: defaultdict[str, list[float]](list) + ) + cui_giou: defaultdict[str, list[float]] = Field( + default_factory=lambda: defaultdict[str, list[float]](list) + ) + cui_cohen_k: defaultdict[str, list[float]] = Field( + default_factory=lambda: defaultdict[str, list[float]](list) + ) + + +class OverallMetrics(BaseModel): + """Project / Mode level metrics, calculated from RawStats.""" + precision: float = 0.0 + recall: float = 0.0 + f1: float = 0.0 + + # Number of labels where it is not possible to generate an entity + no_tokens: int = 0 + # Number of labels in entire project where it is not + # possible to generate an entity + no_tokens_ratio: float = 0.0 + + tp: int = 0 + fp: int = 0 + fn: int = 0 + + char_iou: float = 0.0 + char_giou: float = 0.0 + char_cohen_k: float = 0.0 + + +class CUIMetrics(BaseModel): + """Metrics on a per cui basis.""" + precision: float = 0.0 + recall: float = 0.0 + f1: float = 0.0 + + tp: int = 0 + fp: int = 0 + fn: int = 0 + + char_iou: float = 0.0 + char_giou: float = 0.0 + char_cohen_k: float = 0.0 + + char_iou_n: float = 0.0 + char_giou_n: float = 0.0 + char_cohen_k_n: float = 0.0 + + +class Metrics(BaseModel): + """Calculated metrics for a single evaluation mode.""" + + overall: OverallMetrics + per_cui: dict[str, CUIMetrics] = Field(default_factory=dict) + + +class GoldAnnotation(TypedDict): + """Validated gold annotation payload after CUI filtering.""" + + start: int + end: int + cuis: list[str] + cui: str + text: str + raw: object + + +class PredictedAnnotation(TypedDict): + """Predicted entity payload used for scoring and metrics.""" + + start: int + end: int + cui: str + text: str + confidence: float + raw: MutableEntity + no_tokens: int + + +class ModeStats(BaseModel): + """Accumulated state and calculated metrics for one evaluation mode.""" + + stats: RawStats = Field(default_factory=RawStats) + metrics: Metrics | None = None + +class ProjectStats(BaseModel): + """Accumulated state and calculated metrics for one project + or all projects.""" + full_pipeline: ModeStats = Field( + default_factory=ModeStats + ) + ner: ModeStats | None = None + linking: ModeStats | None = None + + def get_mode(self, mode: MetricMode) -> ModeStats | None: + normalized = MetricMode(mode) + return getattr(self, normalized.value) + + @classmethod + def create( + cls, + ner: bool = False, + linking: bool = False, + ) -> "ProjectStats": + return cls( + ner=ModeStats() if ner else None, + linking=ModeStats() if linking else None, + ) + +class StatsCollection(BaseModel): + """Accumulated state and calculated metrics for all projects.""" + all_projects: ProjectStats = Field( + default_factory=ProjectStats + ) + projects: dict[int, ProjectStats] = Field( + default_factory=dict + ) + + def get_project_stats(self, project_index: int) -> ProjectStats: + """Return statistics for a single project.""" + return self.projects[project_index] + + def get_aggregate_stats(self) -> ProjectStats: + """Return the aggregate statistics across all projects.""" + return self.all_projects + + @classmethod + def create( + cls, + num_projects: int, + ner: bool = False, + linking: bool = False, + ) -> "StatsCollection": + return cls( + all_projects=ProjectStats.create( + ner=ner, + linking=linking, + ), + projects={ + i: ProjectStats.create( + ner=ner, + linking=linking, + ) + for i in range(num_projects) + }, + ) + +class StatsCalculator: + """Calculates statistics for entity linking.""" + + BUCKET_FULL = MetricMode.FULL + BUCKET_NER = MetricMode.NER + BUCKET_LINKING = MetricMode.LINKING def __init__(self, filters: LinkingFilters, - addl_info: dict, - doc_getter: Callable[[str], Optional[MutableDocument]], cui2info: dict[str, CUIInfo], - use_project_filters: bool = False, - use_overlaps: bool = False, - # use_cui_doc_limit: bool = False, - # use_groups: bool = False, - extra_cui_filter: Optional[set[str]] = None) -> None: + num_projects: int, + ner_performance: bool = False, + linking_performance: bool = False, + ) -> None: self.filters = filters - self.addl_info = addl_info - self.doc_getter = doc_getter self.cui2info = cui2info - self.use_project_filters = use_project_filters - self.use_overlaps = use_overlaps - # self.use_cui_doc_limit = use_cui_doc_limit - # self.use_groups = use_groups - self.extra_cui_filter = extra_cui_filter - self._reset_stats() - - def _reset_stats(self): - self.tp = 0 - self.fp = 0 - self.fn = 0 - self.fps: dict[str, int] = {} - self.fns: dict[str, int] = {} - self.tps: dict[str, int] = {} - self.cui_prec: dict[str, float] = {} - self.cui_rec: dict[str, float] = {} - self.cui_f1: dict[str, float] = {} - self.cui_counts: dict[str, int] = {} - self.examples: dict = {'fp': {}, 'fn': {}, 'tp': {}} - self.fp_docs: set = set() - self.fn_docs: set = set() - - def process_project(self, project: MedCATTrainerExportProject) -> None: - """Process the project. - - This processes each document in the project. + self.reset(num_projects, + ner_performance, + linking_performance) - Args: - project (MedCATTrainerExportProject): The trainer export project. + def reset(self, + num_projects: int, + ner_performance: bool = False, + linking_performance: bool = False) -> None: + self.ner_performance = ner_performance + self.linking_performance = linking_performance + self.num_projects = num_projects + self.stats = StatsCollection().create( + num_projects=self.num_projects, + ner=self.ner_performance, + linking=self.linking_performance + ) + + def _extract_gold_annotations( + self, + doc: MedCATTrainerExportDocument + ) -> list[GoldAnnotation]: + """Extract validated gold annotations, supporting multi-CUI options.""" + gold_anns: list[GoldAnnotation] = [] + + for ann in doc['annotations']: + if not ann.get('validated', True): + continue + if ann.get('killed', False) or ann.get('deleted', False): + continue + + # Support both single CUI and multiple acceptable CUIs. + acceptable_cuis = ann.get('acceptable_cuis', ann['cui']) + if isinstance(acceptable_cuis, list): + cuis = acceptable_cuis + else: + cuis = [acceptable_cuis] + + # Filter to valid CUIs. + valid_cuis: list[str] = [ + cui + for cui in cuis + if isinstance(cui, str) + and self.filters.check_filters(cui) + ] + if valid_cuis: + gold_anns.append({ + 'start': ann['start'], + 'end': ann['end'], + 'cuis': valid_cuis, # List of acceptable CUIs + 'cui': valid_cuis[0], # For counting + 'text': ann['value'], + 'raw': ann + }) + return gold_anns + + def _extract_predictions( + self, + predictions: list[MutableEntity], + apply_filters: bool = True, + ) -> list[PredictedAnnotation]: + """Extract relevant info from predicted entities.""" + extracted: list[PredictedAnnotation] = [] + + for ent in predictions: + if apply_filters and not self.filters.check_filters(ent.cui): + continue + + extracted.append({ + 'start': ent.base.start_char_index, + 'end': ent.base.end_char_index, + 'cui': ent.cui, + 'text': ent.base.text, + 'confidence': float(ent.context_similarity), + 'raw': ent, + 'no_tokens': 1 if ent.id == UNTOKENIZABLE_ENTITY_ID else 0, + }) + + return extracted + + def _count_gold_annotations( + self, + gold_anns: list[GoldAnnotation], + project_index: int, + mode: MetricMode + ) -> None: + """Count gold annotations for a project and all-projects aggregate.""" + project_stats = self.stats.get_project_stats(project_index) + aggregate_stats = self.stats.get_aggregate_stats() + for project_stats in (project_stats, aggregate_stats): + mode_stats = project_stats.get_mode(mode) + if mode_stats is None: + continue + state = mode_stats.stats + if mode == self.BUCKET_NER: + key = "__NER__" + state.cui_gold_counts[key] = ( + state.cui_gold_counts.get(key, 0) + + len(gold_anns) + ) + continue + for gold in gold_anns: + cui = gold["cui"] + state.cui_gold_counts[cui] = ( + state.cui_gold_counts.get(cui, 0) + + 1 + ) + + def _record_tp(self, + state: RawStats, + gold: GoldAnnotation, + pred: PredictedAnnotation) -> None: + """Record a true positive.""" + cui = pred['cui'] + state.tp += 1 + state.cui_tp[cui] = state.cui_tp.get(cui, 0) + 1 + + if cui not in state.examples['tp']: + state.examples['tp'][cui] = [] + state.examples['tp'][cui].append({ + 'gold_text': gold['text'], + 'pred_text': pred['text'], + 'cui': cui, + 'start': pred['start'], + 'confidence': pred['confidence'] + }) + + def _record_fn(self, state: RawStats, gold: GoldAnnotation) -> None: + """Record a false negative.""" + cui = gold['cui'] + state.fn += 1 + state.cui_fn[cui] = state.cui_fn.get(cui, 0) + 1 + + if cui not in state.examples['fn']: + state.examples['fn'][cui] = [] + state.examples['fn'][cui].append({ + 'text': gold['text'], + 'acceptable_cuis': gold['cuis'], + 'start': gold['start'] + }) + + def _record_fp(self, state: RawStats, pred: PredictedAnnotation) -> None: + """Record a false positive.""" + cui = pred['cui'] + state.fp += 1 + state.cui_fp[cui] = state.cui_fp.get(cui, 0) + 1 + + if cui not in state.examples['fp']: + state.examples['fp'][cui] = [] + state.examples['fp'][cui].append({ + 'text': pred['text'], + 'cui': cui, + 'start': pred['start'], + 'confidence': pred['confidence'] + }) + + def _record_no_tokens(self, state: RawStats, pred: PredictedAnnotation) -> None: + """Record a prediction with no tokens (ID -1000).""" + # When there's an entity with no way for the tokenizer to parse it + # (commonly, this means that it's a subtoken span i.e. mRBC -> RBC isn't viable) + # There's no tokens, throwing an error at get_tokens. + # Treat it as a gold-like false negative so the recorded payload matches + # the expected annotation schema used by the rest of the scorer. + gold: GoldAnnotation = { + 'start': pred['start'], + 'end': pred['end'], + 'cuis': [pred['cui']], + 'cui': pred['cui'], + 'text': pred['text'], + 'raw': pred['raw'], + } + cui = gold['cui'] + state.no_tokens += 1 + state.cui_no_tokens[cui] = state.cui_no_tokens.get(cui, 0) + 1 + self._record_fn(state, gold) + + def _find_matching_prediction( + self, + gold: GoldAnnotation, + predictions: list[PredictedAnnotation], + matched_preds: set[int] + ) -> int | None: """ - project_name = cast(str, project.get('name')) - project_id = cast(str, project.get('id')) - - documents = project["documents"] - for dind, doc in tqdm( - enumerate(documents), - desc="Stats document", - total=len(documents), - leave=False, - ): - self.process_document(project_name, project_id, doc) - - def process_document(self, project_name: str, project_id: str, - doc: MedCATTrainerExportDocument - ) -> None: - """Process the trainer export document. + Find a prediction that matches this gold annotation. - Args: - project_name (str): The project within which this document lies. - project_id (str): The project ID for the project. - doc (MedCATTrainerExportDocument): The trainer export document. + Matching criteria: + - Same start position (can be relaxed for fuzzy matching) + - Predicted CUI is in gold's acceptable CUIs + - Not already matched """ - anns = doc['annotations'] - - # Apply document level filtering, in this case project_filter is - # ignored while the extra_cui_filter is respected still - # if self.use_cui_doc_limit: - # _cuis = set([ann['cui'] for ann in anns]) - # if _cuis: - # self.filters.cuis = intersect_nonempty_set(_cuis, - # self.extra_cui_filter) - # else: - # self.filters.cuis = {'empty'} - - mut_doc: MutableDocument = self.doc_getter( - doc['text']) # type: ignore - - p_anns = mut_doc.linked_ents # or all ents? - - (anns_norm, anns_norm_neg, - anns_examples, _) = self._preprocess_annotations( - project_name, project_id, doc, anns) - - p_anns_norm, p_anns_examples = self._process_p_anns( - project_name, project_id, doc, p_anns) - self._count_p_anns_norm(doc, anns_norm, anns_norm_neg, - p_anns_norm, p_anns_examples) - self._process_anns_norm(doc, anns_norm, p_anns_norm, anns_examples) - - def _process_anns_norm(self, doc: MedCATTrainerExportDocument, - anns_norm: list[tuple[int, str]], - p_anns_norm: list[tuple[int, str]], - anns_examples: list[dict]) -> None: - for iann, ann in enumerate(anns_norm): - if ann not in p_anns_norm: - cui = ann[1] - self.fn += 1 - self.fn_docs.add(doc.get('name', 'unk')) - - self.fns[cui] = self.fns.get(cui, 0) + 1 - examples = self.examples['fn'].get(cui, []) - self.examples['fn'][cui] = examples + [anns_examples[iann]] - - def _process_p_anns(self, project_name: str, project_id: str, - doc: MedCATTrainerExportDocument, - p_anns: list[MutableEntity] - ) -> tuple[list[tuple[int, str]], list[dict]]: - p_anns_norm: list[tuple[int, str]] = [] - p_anns_examples: list[dict] = [] - for ann in p_anns: - cui = ann.cui - - p_anns_norm.append((ann.base.start_char_index, cui)) - p_anns_examples.append(self._create_annotation_2( - project_name, project_id, cui, doc, ann)) - return p_anns_norm, p_anns_examples - - def _count_p_anns_norm(self, doc: MedCATTrainerExportDocument, - anns_norm: list[tuple[int, str]], - anns_norm_neg: list[tuple[int, str]], - p_anns_norm: list[tuple[int, str]], - p_anns_examples: list[dict]) -> None: - for iann, ann in enumerate(p_anns_norm): - cui = ann[1] - if ann in anns_norm: - self.tp += 1 - self.tps[cui] = self.tps.get(cui, 0) + 1 - - example = p_anns_examples[iann] - - examples = self.examples['tp'].get(cui, []) - self.examples['tp'][cui] = examples + [example] + for idx, pred in enumerate(predictions): + if idx in matched_preds: + continue + + # Exact span match + if pred['start'] == gold['start']: + # Check if predicted CUI is acceptable + if pred['cui'] in gold['cuis']: + return idx + + return None + + def _score_annotations(self, + gold_anns: list[GoldAnnotation], + pred_anns: list[PredictedAnnotation], + project_index: int, + mode: MetricMode, + filter_fp_by_cui: bool = True) -> None: + # Track which predictions have been matched + matched_preds: set[int] = set() + aggregate_stats = self.stats.get_aggregate_stats() + project_stats = self.stats.get_project_stats(project_index) + all_projects_state = aggregate_stats.get_mode(mode) + project_state = project_stats.get_mode(mode) + + if all_projects_state is None or project_state is None: + return + + # this is a bit counter intuitive. + # essentially if you're looking at the linking performance, + # then there maybe entities with no tokens (due to spacy i.e + # [m'RNA'] not being representated) So you have to check the + # ner'd spans for linking performance. + if mode == self.BUCKET_LINKING: + for pred in pred_anns: + if pred['no_tokens'] == 1: + self._record_no_tokens(all_projects_state.stats, pred) + self._record_no_tokens(project_state.stats, pred) + + # NOTE: All predictions where ID are -1000 are false positives. + # this should only really happen on the linker testing, as it's a perfect + # NER step which is trying to create tokenless entities. + # Phase 1: Match gold annotations to predictions (find TPs and FNs) + for gold in gold_anns: + if not gold['cuis']: + # No valid CUIs for this gold annotation, skip it + continue + match_idx = self._find_matching_prediction( + gold, pred_anns, matched_preds) + + if match_idx is not None: + # True Positive + matched_preds.add(match_idx) + pred = pred_anns[match_idx] + self._record_tp(all_projects_state.stats, gold, pred) + self._record_tp(project_state.stats, gold, pred) else: - self.fp += 1 - self.fps[cui] = self.fps.get(cui, 0) + 1 - self.fp_docs.add(doc.get('name', 'unk')) - - # Add example for this FP prediction - example = p_anns_examples[iann] - if ann in anns_norm_neg: - # Means that it really was annotated as negative - example['real_fp'] = True - - examples = self.examples['fp'].get(cui, []) - self.examples['fp'][cui] = examples + [example] - - def _create_annotation(self, project_name: str, project_id: str, cui: str, - doc: MedCATTrainerExportDocument, - ann: MedCATTrainerExportAnnotation) -> dict: - return {"text": doc['text'][max(0, ann['start'] - 60):ann['end'] + 60], - "cui": cui, - "start": ann['start'], - "end": ann['end'], - "source value": ann['value'], - "acc": 1, - "project name": project_name, - "document name": doc.get('name'), - "project id": project_id, - "document id": doc.get('id')} - - def _create_annotation_2(self, project_name: str, project_id: str, - cui: str, doc: MedCATTrainerExportDocument, - ann: MutableEntity) -> dict: - start = max(0, ann.base.start_char_index - 60) - end = ann.base.end_char_index + 60 - return {"text": doc['text'][start:end], - "cui": cui, - "start": ann.base.start_char_index, - "end": ann.base.start_char_index, - "source value": ann.base.text, - "acc": float(ann.context_similarity), - "project name": project_name, - "document name": doc.get('name'), - "project id": project_id, - "document id": doc.get('id')} - - def _preprocess_annotations(self, project_name: str, project_id: str, - doc: MedCATTrainerExportDocument, - anns: list[MedCATTrainerExportAnnotation] - ) -> tuple[list[tuple[int, str]], - list[tuple[int, str]], - list[dict], - list[str]]: - anns_norm: list[tuple[int, str]] = [] - anns_norm_neg: list[tuple[int, str]] = [] - anns_examples: list[dict] = [] - anns_norm_cui: list[str] = [] + # False Negative + self._record_fn(all_projects_state.stats, gold) + self._record_fn(project_state.stats, gold) + + # Phase 2: Remaining predictions are False Positives + for idx, pred in enumerate(pred_anns): + if idx in matched_preds: + continue + if filter_fp_by_cui and not self.filters.check_filters(pred['cui']): + continue + self._record_fp(all_projects_state.stats, pred) + self._record_fp(project_state.stats, pred) + + def _to_ner_views(self, + gold_anns: list[GoldAnnotation], + pred_anns: list[PredictedAnnotation] + ) -> tuple[list[GoldAnnotation], list[PredictedAnnotation]]: + ner_cui = '__NER__' + eval_pred_anns: list[PredictedAnnotation] = [ + {**pred, "cui": ner_cui} + for pred in pred_anns + ] + + eval_gold_anns: list[GoldAnnotation] = [ + {**gold, "cuis": [ner_cui], "cui": ner_cui} + for gold in gold_anns + ] + return eval_gold_anns, eval_pred_anns + + def _build_character_sets( + self, + anns: list[PredictedAnnotation] | list[GoldAnnotation], + ) -> dict[str, set[int]]: + chars_by_cui = defaultdict(set) + for ann in anns: + start = int(ann['start']) + end = int(ann['end']) cui = ann['cui'] - if self.filters.check_filters(cui): - - if (ann.get('validated', True) and - (not ann.get('killed', False) and not ann.get('deleted', - False))): - anns_norm.append((ann['start'], cui)) - anns_examples.append(self._create_annotation( - project_name, project_id, cui, doc, ann)) - elif (ann.get('validated', True) and - (ann.get('killed', False) or ann.get('deleted', False))): - anns_norm_neg.append((ann['start'], cui)) - - if ann.get("validated", True): - # This is used to test was someone annotating for this - # CUI in this document - anns_norm_cui.append(cui) - self.cui_counts[cui] = self.cui_counts.get(cui, 0) + 1 - return anns_norm, anns_norm_neg, anns_examples, anns_norm_cui - - def finalise_report(self, epoch: int, do_print: bool = True): - """Finalise the report / metrics. + char_idxs = set(range(start, end)) + chars_by_cui[cui].update(char_idxs) - This prints out the overall metrics and calculates per CUI metrics. + + return dict(chars_by_cui) + + def _character_cohen_kappa( + self, + gold_chars: set[int], + pred_chars: set[int], + document_length: int, + ) -> float: + """ + The voices in my chatbot told me this is faster than the + sklearn implementation, and it is also more memory efficient. + + Testing shows same metrics, and halving computation speed. + """ + + tp = len(gold_chars & pred_chars) + fp = len(pred_chars - gold_chars) + fn = len(gold_chars - pred_chars) + tn = document_length - tp - fp - fn + + total = document_length + + if total == 0: + return 1.0 + + # Observed agreement + po = (tp + tn) / total + + # Expected agreement + gold_positive = tp + fn + gold_negative = fp + tn + + pred_positive = tp + fp + pred_negative = fn + tn + + pe = ( + (gold_positive * pred_positive) + + + (gold_negative * pred_negative) + ) / (total * total) + + denominator = 1 - pe + + if denominator == 0: + # Perfect agreement or no variation + return 1.0 + + return (po - pe) / denominator + + def _calculate_document_character_scores( + self, + gold_anns: list[GoldAnnotation], + pred_anns: list[PredictedAnnotation], + doc_length: int, + ) -> tuple[dict[str, float], dict[str, float], dict[str, float]]: + """Compute per-CUI character score components for one document.""" + gold_chars_by_cui = self._build_character_sets(gold_anns) + pred_chars_by_cui = self._build_character_sets(pred_anns) + + # For standard IoU and Cohen's Kappa: + # include CUIs appearing in either gold or prediction. + all_cuis = set(gold_chars_by_cui) | set(pred_chars_by_cui) + + # For GIoU: Gold Label Intersection over Union, + # we only evaluate CUIs that are present in labels. + # only include CUIs present in gold. + gold_cuis = set(gold_chars_by_cui) + + per_cui_ious: dict[str, float] = {} + per_cui_gious: dict[str, float] = {} + per_cui_kappas: dict[str, float] = {} + + for cui in all_cuis: + gold_chars = gold_chars_by_cui.get(cui, set()) + pred_chars = pred_chars_by_cui.get(cui, set()) + + intersection = gold_chars & pred_chars + union = gold_chars | pred_chars + + # Character IoU + iou = len(intersection) / len(union) if union else 1.0 + per_cui_ious[cui] = iou + + # Gold IoU / GIoU + # Only evaluated for CUIs present in gold. + # Prediction-only CUIs are ignored. + if cui in gold_cuis: + giou = len(intersection) / len(gold_chars) + per_cui_gious[cui] = giou + + per_cui_kappas[cui] = self._character_cohen_kappa( + gold_chars, + pred_chars, + doc_length, + ) + + return per_cui_ious, per_cui_gious, per_cui_kappas + + def _update_project_stats( + self, + project_state: ModeStats, + all_project_state: ModeStats, + per_cui_ious: dict[str, float], + per_cui_gious: dict[str, float], + per_cui_kappas: dict[str, float], + ) -> None: + """Apply a document's character metric values to the project state.""" + for cui, iou in per_cui_ious.items(): + project_state.stats.cui_iou[cui].append(iou) + all_project_state.stats.cui_iou[cui].append(iou) + + for cui, giou in per_cui_gious.items(): + project_state.stats.cui_giou[cui].append(giou) + all_project_state.stats.cui_giou[cui].append(giou) + + for cui, kappa in per_cui_kappas.items(): + project_state.stats.cui_cohen_k[cui].append(kappa) + all_project_state.stats.cui_cohen_k[cui].append(kappa) + + # Average the per-CUI IoUs rather than merging character sets. + # This preserves CUI identity. + doc_iou = ( + sum(per_cui_ious.values()) / len(per_cui_ious) + if per_cui_ious else 0.0 + ) + + # Only gold CUIs contribute to GIoU, so we average over those. + doc_giou = ( + sum(per_cui_gious.values()) / len(per_cui_gious) + if per_cui_gious else 0.0 + ) + + # cohen's kappa is averaged over all CUIs, including those only in predictions. + doc_cohen_k = ( + sum(per_cui_kappas.values()) / len(per_cui_kappas) + if per_cui_kappas else 1.0 + ) + + project_state.stats.iou_sum += doc_iou + all_project_state.stats.iou_sum += doc_iou + + project_state.stats.giou_sum += doc_giou + all_project_state.stats.giou_sum += doc_giou + + project_state.stats.cohen_k_sum += doc_cohen_k + all_project_state.stats.cohen_k_sum += doc_cohen_k + + project_state.stats.char_docs += 1 + all_project_state.stats.char_docs += 1 + + def _score_character_annotations(self, + gold_anns: list[GoldAnnotation], + pred_anns: list[PredictedAnnotation], + project_index: int, + mode: MetricMode, + doc_length: int) -> None: + """ + Calculate: + - Character Intersection over Union (IoU) for gold and predicted annotations. + - Gold label Character Intersection over Union (IoU) for gold and + predicted annotations. + - Cohen's Kappa for gold and predicted annotations. + + Cheat sheet of what we're generating: + # iou = sum of document-level macro IoUs + # -> divide by number of documents + # giou = sum of document-level macro GIoUs + # -> divide by number of documents + # cohen_k = sum of document-level macro Kappas + # -> divide by number of documents + # cui_iou[CUI] = sum of per-document IoU for that CUI + # -> divide by number of documents containing that CUI + # cui_giou[CUI] = sum of per-document GIoU for that CUI + # -> divide by number of documents containing that CUI in gold + # cui_cohen_k[CUI] = sum of per-document CUI-specific Kappa + # -> divide by number of documents where the CUI is evaluated + """ + aggregate_stats = self.stats.get_aggregate_stats() + project_stats = self.stats.get_project_stats(project_index) + all_project_state = aggregate_stats.get_mode(mode) + project_state = project_stats.get_mode(mode) + + if project_state is None or all_project_state is None: + return + + per_cui_ious, per_cui_gious, per_cui_kappas = ( + self._calculate_document_character_scores( + gold_anns, + pred_anns, + doc_length, + ) + ) + + self._update_project_stats( + project_state, + all_project_state, + per_cui_ious, + per_cui_gious, + per_cui_kappas, + ) + + def process_document( + self, + doc: MedCATTrainerExportDocument, + project_index: int, + predictions: list[MutableEntity], + mode: MetricMode, + calculate_ner_performance: bool = False, + ) -> None: + """ + Process a single document's annotations and predictions. Args: - epoch (int): The number of the current epoch. - do_print (bool, optional): Whether to print the output. - Defaults to True. + doc: Gold-standard annotated document + predictions: Model's predicted entities """ - try: - if self.tp + self.fp == 0: - prec = 0.0 - else: - prec = self.tp / (self.tp + self.fp) - if self.tp + self.fp == 0: - rec = 0.0 - else: - rec = self.tp / (self.tp + self.fn) - if prec == 0 and rec == 0: - f1 = 0.0 - else: - f1 = 2 * (prec * rec) / (prec + rec) - if do_print: - print("Epoch: {}, Prec: {}, Rec: {}, F1: {}\n".format( - epoch, prec, rec, f1)) - print("Docs with false positives: {}\n".format("; ".join( - [str(x) for x in list(self.fp_docs)[0:10]]))) - print("Docs with false negatives: {}\n".format("; ".join( - [str(x) for x in list(self.fn_docs)[0:10]]))) - - # Sort fns & prec - fps = {k: v for k, v in sorted(self.fps.items(), - key=lambda item: item[1], reverse=True)} - fns = {k: v for k, v in sorted(self.fns.items(), - key=lambda item: item[1], reverse=True)} - tps = {k: v for k, v in sorted(self.tps.items(), - key=lambda item: item[1], reverse=True)} - - # F1 per concept - for cui in tps.keys(): - prec = tps[cui] / (tps.get(cui, 0) + fps.get(cui, 0)) - rec = tps[cui] / (tps.get(cui, 0) + fns.get(cui, 0)) - f1 = 2 * (prec * rec) / (prec + rec) - self.cui_prec[cui] = prec - self.cui_rec[cui] = rec - self.cui_f1[cui] = f1 - - # Get top 10 - pr_fps = [(self._get_pref_name(cui), - cui, fps[cui]) for cui in list(fps.keys())[0:10]] - pr_fns = [(self._get_pref_name(cui), - cui, fns[cui]) for cui in list(fns.keys())[0:10]] - pr_tps = [(self._get_pref_name(cui), - cui, tps[cui]) for cui in list(tps.keys())[0:10]] - - if do_print: - print("\n\nFalse Positives\n") - for one in pr_fps: - print("{:70} - {:20} - {:10}".format(str(one[0])[0:69], - str(one[1])[0:19], - one[2])) - print("\n\nFalse Negatives\n") - for one in pr_fns: - print("{:70} - {:20} - {:10}".format(str(one[0])[0:69], - str(one[1])[0:19], - one[2])) - print("\n\nTrue Positives\n") - for one in pr_tps: - print("{:70} - {:20} - {:10}".format(str(one[0])[0:69], - str(one[1])[0:19], - one[2])) - print("*" * 110 + "\n") - - except Exception: - traceback.print_exc() + full_pipe_gold_anns = self._extract_gold_annotations(doc) + full_pipe_pred_anns = self._extract_predictions(predictions) + + self._count_gold_annotations(full_pipe_gold_anns, project_index, mode) + self._score_annotations( + full_pipe_gold_anns, + full_pipe_pred_anns, + project_index, + mode=mode, + filter_fp_by_cui=True + ) + self._score_character_annotations( + full_pipe_gold_anns, + full_pipe_pred_anns, + project_index, + mode=mode, doc_length=len(doc['text']) + ) + + # This gets called in the full pipeline call, if ner performance is called. + if calculate_ner_performance: + ner_gold_anns, ner_pred_anns = self._to_ner_views( + full_pipe_gold_anns, full_pipe_pred_anns) + self._count_gold_annotations(ner_gold_anns, project_index, + mode=self.BUCKET_NER) + self._score_annotations(ner_gold_anns, ner_pred_anns, + project_index, mode=self.BUCKET_NER, + filter_fp_by_cui=False) + self._score_character_annotations( + ner_gold_anns, + ner_pred_anns, + project_index, + mode=self.BUCKET_NER, + doc_length=len(doc['text']) + ) + + def process_project(self, project: MedCATTrainerExportProject, + project_index: int, + entity_getter: Callable[[str], list[MutableEntity]], + mode: MetricMode, + calculate_ner_performance: bool = False, + use_project_filters: bool = False, + extra_cui_filter: set[str] | None = None + ) -> None: + """Process all documents in a project. + + Args: + project: The project data containing documents and annotations. + project_index: Index of the project in the export. + entity_getter: Function to get predicted entities from text. + mode: Evaluation mode (full, ner, linking). + calculate_ner_performance: Whether to calculate NER performance. + use_project_filters: Whether to apply project-specific filters. + extra_cui_filter: Additional CUI filter to apply. + """ + with project_filters(self.filters, + project, + extra_cui_filter, + use_project_filters): + for doc in tqdm(project['documents'], + desc='Documents'): + predictions = entity_getter(doc['text']) + self.process_document( + doc, + project_index, + predictions, + mode=mode, + calculate_ner_performance=calculate_ner_performance, + ) + + def _get_linked_ents(self, cat: CAT, text: str) -> list[MutableEntity]: + """Required for mypy cleanliness""" + doc = cat(text) + if doc is None: + return [] + return doc.linked_ents + + def process_export(self, cat: CAT, export: MedCATTrainerExport, + mode: MetricMode, + calculate_ner_performance: bool = False, + use_project_filters: bool = False, + extra_cui_filter: set[str] | None = None, + filter_before_disamb: bool = False) -> None: + """Process all projects in the export. + + Args: + cat: The MedCAT CAT instance for entity linking. + export: The MedCAT trainer export data. + mode: Evaluation mode (full, ner, linking). + calculate_ner_performance: Whether to calculate NER performance. + use_project_filters: Whether to apply project-specific filters. + extra_cui_filter: Additional CUI filter to apply. + filter_before_disamb: Whether to filter entities before disambiguation. + """ + if filter_before_disamb: + cat.config.components.linking.filter_before_disamb = True + for i, proj in tqdm(enumerate(export['projects']), desc='Projects'): + self.process_project( + proj, + i, + lambda text: self._get_linked_ents(cat, text), + mode=mode, + calculate_ner_performance=calculate_ner_performance, + use_project_filters=use_project_filters, + extra_cui_filter=extra_cui_filter + ) + + @staticmethod + def _compute_prf(tp: int, fp: int, fn: int, no_tokens: int) -> dict: + """Compute precision, recall, F1.""" + prec = tp / (tp + fp) if (tp + fp) > 0 else 0.0 + rec = tp / (tp + fn) if (tp + fn) > 0 else 0.0 + f1 = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0.0 + no_tokens_ratio = no_tokens / (tp + fn) if (tp + fn) > 0 else 0.0 + return { + 'precision': prec, + 'recall': rec, + 'f1': f1, + 'no_tokens': no_tokens, + 'no_tokens_ratio': f'{no_tokens_ratio:.4f}' + } + + def _get_cui_name(self, cui: str) -> str: + """Get preferred name for CUI.""" + info = self.cui2info.get(cui) + if info: + return info.get('preferred_name') or list(info['names'])[0] + return cui + + def _safe_mean(self, values): + return sum(values) / len(values) if values else 0.0 + + def _prepare_metrics(self, + raw_stats: RawStats) -> tuple[OverallMetrics, dict[str, dict]]: + """Prepare overall and per-CUI metrics from raw accumulated state.""" + # project metrics + prf_values = self._compute_prf( + raw_stats.tp, + raw_stats.fp, + raw_stats.fn, + raw_stats.no_tokens, + ) + + if raw_stats.char_docs > 0: + char_iou = raw_stats.iou_sum / raw_stats.char_docs + char_giou = raw_stats.giou_sum / raw_stats.char_docs + char_cohen_k = raw_stats.cohen_k_sum / raw_stats.char_docs + else: + char_iou = 0.0 + char_giou = 0.0 + char_cohen_k = 0.0 + + overall = OverallMetrics( + precision=prf_values['precision'], + recall=prf_values['recall'], + f1=prf_values['f1'], + no_tokens=raw_stats.no_tokens, + no_tokens_ratio=float(prf_values['no_tokens_ratio']), + tp=raw_stats.tp, + fp=raw_stats.fp, + fn=raw_stats.fn, + char_iou=char_iou, + char_giou=char_giou, + char_cohen_k=char_cohen_k, + ) + + # cui metrics + all_cuis = ( + set(raw_stats.cui_tp) + | set(raw_stats.cui_fp) + | set(raw_stats.cui_fn) + | set(raw_stats.cui_iou) + | set(raw_stats.cui_giou) + | set(raw_stats.cui_cohen_k) + ) + + per_cui: dict[str, dict] = {} + + for cui in all_cuis: + tp = raw_stats.cui_tp.get(cui, 0) + fp = raw_stats.cui_fp.get(cui, 0) + fn = raw_stats.cui_fn.get(cui, 0) + no_tokens = raw_stats.cui_no_tokens.get(cui, 0) + + cui_iou_scores = raw_stats.cui_iou.get(cui, []) + cui_giou_scores = raw_stats.cui_giou.get(cui, []) + cui_k_scores = raw_stats.cui_cohen_k.get(cui, []) + + per_cui[cui] = { + "name": self._get_cui_name(cui), + **self._compute_prf( + tp, + fp, + fn, + no_tokens, + ), + "tp": tp, + "fp": fp, + "fn": fn, + "char_iou": self._safe_mean(cui_iou_scores), + "char_giou": self._safe_mean(cui_giou_scores), + "char_cohen_k": self._safe_mean(cui_k_scores), + "char_iou_n": len(cui_iou_scores), + "char_giou_n": len(cui_giou_scores), + "char_cohen_k_n": len(cui_k_scores), + } + + return overall, per_cui + def compute_metrics( + self, + stats: ProjectStats, + mode: MetricMode + ) -> None: + """Compute overall and per-CUI metrics for a given mode.""" + + mode_stats = stats.get_mode(mode) + if mode_stats is None: + return + overall, per_cui = self._prepare_metrics(mode_stats.stats) + + # Store computed metrics in the ModeStats object + mode_stats.metrics = Metrics( + overall=overall, + per_cui={ + cui: CUIMetrics(**metrics) + for cui, metrics in per_cui.items() + }, + ) + + def compute_all_metrics(self, + ner_performance: bool = True, + linking_performance: bool = True) -> None: + """Compute metrics for all projects and the aggregate.""" + stats = self.stats.get_aggregate_stats() + self.compute_metrics(stats, StatsCalculator.BUCKET_FULL) + if ner_performance: + self.compute_metrics(stats, StatsCalculator.BUCKET_NER) + if linking_performance: + self.compute_metrics(stats, StatsCalculator.BUCKET_LINKING) + + if self.num_projects > 1: + for i in range(self.num_projects): + stats = self.stats.get_project_stats(i) + self.compute_metrics(stats, + StatsCalculator.BUCKET_FULL) + if ner_performance: + self.compute_metrics(stats, + StatsCalculator.BUCKET_NER) + if linking_performance: + self.compute_metrics(stats, + StatsCalculator.BUCKET_LINKING) + + # these 3 functions are just copied from previous, + # they get nice names for concepts def _empty(self, cui: str) -> CUIInfo: return get_new_cui_info( cui=cui, preferred_name=cui, names=set((cui, ))) @@ -314,130 +989,164 @@ def _get_or_empty(self, cui: str) -> CUIInfo: def _get_pref_name(self, cui: str) -> str: info = self._get_or_empty(cui) return info['preferred_name'] or list(info['names'])[0] + + def print_stats(self, + epoch: int, + mode_stats: ModeStats, + n_samples: int = 10, + stream: TextIO | None = None) -> None: + """Finalise the report / metrics. + + This prints out the overall metrics and calculates per CUI metrics. + + Args: + epoch (int): The number of the current epoch. + mode_stats (ModeStats): The statistics for the current mode. + n_samples (int): Number of entries to print for each section. + stream (TextIO | None): Optional output stream to direct the report + to instead of stdout. + """ + if mode_stats.metrics is None: + raise ValueError( + "Metrics have not been computed yet. " + "Call compute_metrics() first." + ) + print("Epoch: {}, Prec: {}, Rec: {}, F1: {}\n".format( + epoch, + mode_stats.metrics.overall.precision, + mode_stats.metrics.overall.recall, + mode_stats.metrics.overall.f1 + ), file=stream + ) + + # Sort fns & prec + fps = {k: v for k, v in sorted(mode_stats.metrics.per_cui.items(), + key=lambda item: item[1].fp, reverse=True)} + fns = {k: v for k, v in sorted(mode_stats.metrics.per_cui.items(), + key=lambda item: item[1].fn, reverse=True)} + tps = {k: v for k, v in sorted(mode_stats.metrics.per_cui.items(), + key=lambda item: item[1].tp, reverse=True)} - def unwrap(self) -> tuple[ + # Get top 5 + pr_fps = [(self._get_pref_name(cui), + cui, fps[cui]) for cui in list(fps.keys())[0:n_samples]] + pr_fns = [(self._get_pref_name(cui), + cui, fns[cui]) for cui in list(fns.keys())[0:n_samples]] + pr_tps = [(self._get_pref_name(cui), + cui, tps[cui]) for cui in list(tps.keys())[0:n_samples]] + + print("\n\nFalse Positives\n", file=stream) + for one in pr_fps: + print("{:70} - {:20} - {:10}".format( + str(one[0])[0:69], + str(one[1])[0:19], + one[2].fp), + file=stream + ) + print("\n\nFalse Negatives\n", file=stream) + for one in pr_fns: + print("{:70} - {:20} - {:10}".format( + str(one[0])[0:69], + str(one[1])[0:19], + one[2].fn), + file=stream + ) + print("\n\nTrue Positives\n", file=stream) + for one in pr_tps: + print("{:70} - {:20} - {:10}".format( + str(one[0])[0:69], + str(one[1])[0:19], + one[2].tp), + file=stream + ) + print("*" * 110 + "\n", file=stream) + + def legacy_stats(self, mode_stats: "ModeStats") -> tuple[ dict[str, int], dict[str, int], dict[str, int], dict[str, float], dict[str, float], dict[str, float], dict[str, int], dict ]: - return (self.fps, self.fns, self.tps, - self.cui_prec, self.cui_rec, self.cui_f1, - self.cui_counts, self.examples) + per_cui = mode_stats.metrics.per_cui if mode_stats.metrics is not None else {} + to_return = ( + mode_stats.stats.cui_fp, + mode_stats.stats.cui_fn, + mode_stats.stats.cui_tp, + {cui: metrics.precision for cui, metrics in per_cui.items()}, + {cui: metrics.recall for cui, metrics in per_cui.items()}, + {cui: metrics.f1 for cui, metrics in per_cui.items()}, + mode_stats.stats.cui_gold_counts, + mode_stats.stats.examples, + ) + return to_return - @classmethod - def from_cat(cls, cat: CAT, - use_project_filters: bool = False, - use_overlaps: bool = False, - # use_cui_doc_limit: bool = False, - # use_groups: bool = False, - extra_cui_filter: Optional[set[str]] = None - ) -> 'StatsBuilder': - """Get the stats builder from a model pack and some extra information. +def get_stats_calculator(cat: CAT, + data: MedCATTrainerExport, + epoch: int = 0, + use_project_filters: bool = False, + use_overlaps: bool = False, + ner_performance: bool = False, + linking_performance: bool = False, + extra_cui_filter: Optional[set[str]] = None, + do_print: bool = True,) -> StatsCalculator: + calculator = StatsCalculator( + filters=cat.config.components.linking.filters, + cui2info=cat.cdb.cui2info, + num_projects=len(data['projects']), + ner_performance=ner_performance, + linking_performance=linking_performance + ) + # Always compute full pipeline metrics. + # If ner is of interest then also compute NER metrics from the same pass. + calculator.process_export( + cat, + data, + mode=StatsCalculator.BUCKET_FULL, + calculate_ner_performance=ner_performance, + use_project_filters=use_project_filters, + extra_cui_filter=extra_cui_filter, + ) + # Optionally compute linking-only metrics with perfect upstream NER. + if linking_performance: + with dataset_aware_component(cat, CoreComponentType.ner, data): + calculator.process_export( + cat, + data, + mode=StatsCalculator.BUCKET_LINKING, + use_project_filters=use_project_filters, + extra_cui_filter=extra_cui_filter, + ) - Args: - cat (CAT): - The model pack. - use_project_filters (bool, optional): - Whether to use per project filters. Defaults to False. - use_overlaps (bool, optional): - Whether to allow overlaps. Defaults to False. - extra_cui_filter (Optional[set[str]], optional): - Extra CUI filter. Defaults to None. - - Returns: - StatsBuilder: The stats builder. - """ - return StatsBuilder(addl_info=cat.cdb.addl_info, - filters=cat.config.components.linking.filters, - doc_getter=cat.__call__, - # cui2group=cat.cdb.addl_info['cui2group'], - # cui2preferred_name=cat.cdb.cui2preferred_name, - cui2info=cat.cdb.cui2info, - use_project_filters=use_project_filters, - use_overlaps=use_overlaps, - # use_cui_doc_limit=use_cui_doc_limit, - # use_groups=use_groups, - extra_cui_filter=extra_cui_filter) - - -def get_stats(cat: CAT, + calculator.compute_all_metrics(ner_performance, linking_performance) + + if do_print: + to_print = calculator.stats.all_projects.get_mode(StatsCalculator.BUCKET_FULL) + if to_print is None: + raise ValueError("No statistics available for the full pipeline mode.") + calculator.print_stats(epoch, to_print) + return calculator + +def get_stats(cat: CAT, data: MedCATTrainerExport, epoch: int = 0, use_project_filters: bool = False, use_overlaps: bool = False, - # use_cui_doc_limit: bool = False, - # use_groups: bool = False, + ner_performance: bool = False, + linking_performance: bool = False, extra_cui_filter: Optional[set[str]] = None, - do_print: bool = True) -> tuple[ + do_print: bool = True,) -> tuple[ dict[str, int], dict[str, int], dict[str, int], dict[str, float], dict[str, float], dict[str, float], dict[str, int], dict -]: - """TODO: Refactor and make nice - Print metrics on a dataset (F1, P, R), it will also print the concepts - that have the most FP,FN,TP. - - Args: - cat: (CAT): - The model pack. - data (dict): - The json object that we get from MedCATtrainer on export. - epoch (int): - Used during training, so we know what epoch is it. - use_project_filters (bool): - Each project in MedCATtrainer can have filters, do we want to - respect those filters when calculating metrics. - use_overlaps (bool): - Allow overlapping entities, nearly always False as it is very - difficult to annotate overlapping entities. - use_cui_doc_limit (bool): - If True the metrics for a CUI will be only calculated if that CUI - appears in a document, in other words if the document was - annotated for that CUI. Useful in very specific situations when - during the annotation process the set of CUIs changed. - use_groups (bool): - If True concepts that have groups will be combined and stats will - be reported on groups. - extra_cui_filter(Optional[set]): - This filter will be intersected with all other filters, or if all - others are not set then only this one will be used. - do_print (bool): - Whether to print stats out. Defaults to True. - - Returns: - fps (dict): - False positives for each CUI. - fns (dict): - False negatives for each CUI. - tps (dict): - True positives for each CUI. - cui_prec (dict): - Precision for each CUI. - cui_rec (dict): - Recall for each CUI. - cui_f1 (dict): - F1 for each CUI. - cui_counts (dict): - Number of occurrence for each CUI. - examples (dict): - Examples for each of the fp, fn, tp. - Format will be examples['fp']['cui'][]. - """ - builder = StatsBuilder.from_cat(cat, - use_project_filters=use_project_filters, - use_overlaps=use_overlaps, - # use_cui_doc_limit=use_cui_doc_limit, - # use_groups=use_groups, - extra_cui_filter=extra_cui_filter) - for pind, project in tqdm(enumerate(data['projects']), - desc="Stats project", - total=len(data['projects']), - leave=False): - with project_filters(cat.config.components.linking.filters, - project, - builder.extra_cui_filter, - builder.use_project_filters): - builder.process_project(project) - # this is the part that prints out the stats - builder.finalise_report(epoch, do_print=do_print) - return builder.unwrap() + ]: + calculator = get_stats_calculator( + cat=cat, + data=data, + epoch=epoch, + use_project_filters=use_project_filters, + use_overlaps=use_overlaps, + ner_performance=ner_performance, + linking_performance=linking_performance, + extra_cui_filter=extra_cui_filter + ) + full_stats = calculator.stats.all_projects.full_pipeline + return calculator.legacy_stats(full_stats) \ No newline at end of file diff --git a/medcat-v2/medcat/tokenizing/tokens.py b/medcat-v2/medcat/tokenizing/tokens.py index cd7815671..7d7b2f3fe 100644 --- a/medcat-v2/medcat/tokenizing/tokens.py +++ b/medcat-v2/medcat/tokenizing/tokens.py @@ -1,6 +1,10 @@ from typing import Protocol, Optional, Iterator, overload, Any, Type +UNTOKENIZABLE_ENTITY_ID = -1000 +"""Sentinel for entities that could not be mapped to any valid token span.""" + + class BaseToken(Protocol): """Base token protocol. diff --git a/medcat-v2/medcat/utils/training_utils.py b/medcat-v2/medcat/utils/training_utils.py index 1a79efea8..df4824378 100644 --- a/medcat-v2/medcat/utils/training_utils.py +++ b/medcat-v2/medcat/utils/training_utils.py @@ -8,7 +8,9 @@ CoreComponentType, AbstractEntityProvidingComponent) from medcat.config.config import ComponentConfig from medcat.tokenizing.tokenizers import BaseTokenizer -from medcat.tokenizing.tokens import MutableDocument, MutableEntity, MutableToken +from medcat.tokenizing.tokens import ( + MutableDocument, MutableEntity, MutableToken, UNTOKENIZABLE_ENTITY_ID, +) from medcat.data.mctexport import ( MedCATTrainerExport, MedCATTrainerExportDocument, count_all_docs, iter_docs) from medcat.vocab import Vocab @@ -88,9 +90,20 @@ def predict(doc: MutableDocument) -> list[MutableEntity]: anns = _identify_document(doc, dataset)["annotations"] ents: list[MutableEntity] = [] for ann in anns: - tkns = doc.get_tokens(ann["start"], ann["end"]) - # TODO: catch possible exception? - ent = tokens2entity(tkns, doc) + start = ann["start"] + end = ann["end"] + tkns = doc.get_tokens(start, end) + try: + ent = tokens2entity(tkns, doc) + except ValueError: + while not tkns: + # If no tokens found, try expanding the + # range by 1 character on each side + start = max(0, start - 1) + end = end + 1 + tkns = doc.get_tokens(start, end) + ent = tokens2entity(tkns, doc) + ent.id = UNTOKENIZABLE_ENTITY_ID if set_cui: ent.cui = ann["cui"] ents.append(ent) diff --git a/medcat-v2/tests/stats/test_kfold.py b/medcat-v2/tests/stats/test_kfold.py index a67a0ebcb..7a9ac03da 100644 --- a/medcat-v2/tests/stats/test_kfold.py +++ b/medcat-v2/tests/stats/test_kfold.py @@ -215,6 +215,9 @@ def setUpClass(cls) -> None: def setUp(self) -> None: super().setUp() + # return (self.fps, self.fns, self.tps, + # self.cui_prec, self.cui_rec, self.cui_f1, + # self.cui_counts, self.examples) self.reg_stats = reg_stats.get_stats( self.cat, self.mct_export, do_print=False) # TODO - remove diff --git a/medcat-v2/tests/stats/test_stats.py b/medcat-v2/tests/stats/test_stats.py index ed103651a..685419b5a 100644 --- a/medcat-v2/tests/stats/test_stats.py +++ b/medcat-v2/tests/stats/test_stats.py @@ -1,7 +1,13 @@ -from typing import Union +from __future__ import annotations + import os import json - +import re +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Union +from medcat.stats.stats import MetricMode +from medcat.components.types import CoreComponentType from medcat.stats import stats from medcat.data.mctexport import MedCATTrainerExport @@ -12,58 +18,377 @@ os.path.join(os.path.dirname(__file__), "..", "resources")) -class PerfectStatsTests(TrainedModelTests): - PERFECT_STATS_PATH = os.path.join(RESOURCES_PATH, - "mct_export_for_test_exp_perfect.json") +class DummyLinkingFilters: + def __init__(self): + self.cuis = set() + self.cuis_exclude = set() + + def check_filters(self, cui: str) -> bool: + return True + +@dataclass +class DummyToken: + text: str + index: int + start_char_index: int + end_char_index: int + + @property + def base(self): + return SimpleNamespace( + text=self.text, + index=self.index, + start_char_index=self.start_char_index, + end_char_index=self.end_char_index, + ) + + +@dataclass +class DummyEntity: + text: str + start_char_index: int + end_char_index: int + cui: str + context_similarity: float = 1.0 + id: int = 0 + base: SimpleNamespace = field(init=False) + + def __post_init__(self): + self.base = SimpleNamespace( + text=self.text, + start_char_index=self.start_char_index, + end_char_index=self.end_char_index, + index=0, + ) + + +class DummyDocument: + """Minimal object that behaves like a MedCAT document.""" + + def __init__(self, text: str): + self.base = SimpleNamespace(text=text) + self.linked_ents: list[DummyEntity] = [] + + def get_tokens(self, start: int, end: int): + tokens = [] + for idx, match in enumerate(re.finditer(r"\S+", self.base.text)): + s, e = match.span() + if s < end and e > start: + tokens.append( + DummyToken( + text=match.group(), + index=idx, + start_char_index=s, + end_char_index=e, + ) + ) + return tokens + + +class DummyTokenizer: + def entity_from_tokens_in_doc(self, tokens, doc): + if not tokens: + raise ValueError("No tokens to build entity from") + start = tokens[0].base.start_char_index + end = tokens[-1].base.end_char_index + text = doc.base.text[start:end] + return DummyEntity( + text=text, + start_char_index=start, + end_char_index=end, + cui="C0004093", + context_similarity=1.0, + ) + + +# these two are needed for the ner aware performance metrics +class DummyComponent: + def get_type(self): + return CoreComponentType.ner + + +class DummyPipe: + def __init__(self): + self.tokenizer = DummyTokenizer() + self._components = [DummyComponent()] + + def get_component(self, comp_type): + return self._components[0] + + +class DummyCAT: + """Small fake CAT object that supports the stats API.""" + def __init__(self): + self.config = SimpleNamespace( + components=SimpleNamespace( + linking=SimpleNamespace(filters=DummyLinkingFilters()) + ) + ) + self.cdb = SimpleNamespace( + cui2info={ + "195967001": {"preferred_name": "Asthma", "names": {"asthma", "Asthma"}}, + "387458008": {"preferred_name": "Aspirin", "names": {"aspirin", "Aspirin"}}, + } + ) + self.pipe = DummyPipe() + + def __call__(self, text: str): + doc = DummyDocument(text) + for mention, cui in [("asthma", "195967001"), + ("aspirin", "387458008"), + # patient is incorrect linkage + ("patient", "25609006")]: + idx = text.index(mention) + end = idx + len(mention) + ent = DummyEntity( + text=mention, + start_char_index=idx, + end_char_index=end, + cui=cui, + context_similarity=1.0, + ) + doc.linked_ents.append(ent) + return doc + +class DummyCATLinker: + """Small fake CAT object that supports the stats API. Useful for testing the linker.""" + + def __init__(self): + self.config = SimpleNamespace( + components=SimpleNamespace( + linking=SimpleNamespace(filters=DummyLinkingFilters()) + ) + ) + self.cdb = SimpleNamespace( + cui2info={ + "195967001": {"preferred_name": "Asthma", "names": {"asthma", "Asthma"}}, + "387458008": {"preferred_name": "Aspirin", "names": {"aspirin", "Aspirin"}}, + "116154003": {"preferred_name": "Patient", "names": {"patient", "Patient"}}, + "387517004": {"preferred_name": "Paracetamol", "names": {"paracetamol", "Paracetamol"}} + } + ) + self.pipe = DummyPipe() + + def __call__(self, text: str): + doc = DummyDocument(text) + for mention, cui in [("asthma", "195967001"), + ("aspirin", "387458008"), + ("patient", "25609006"), + ("paracetamol", "387517004")]: + idx = text.index(mention) + end = idx + len(mention) + ent = DummyEntity( + text=mention, + start_char_index=idx, + end_char_index=end, + cui=cui, + context_similarity=1.0, + ) + doc.linked_ents.append(ent) + doc.linked_ents.append(ent) + return doc + +def make_fake_test_project() -> dict: + text = "The patient has asthma and takes aspirin, and paracetamol." + annotations = [ + { + "start": text.index("asthma"), + "end": text.index("asthma") + len("asthma"), + "cui": "195967001", + "value": "asthma", + }, + { + "start": text.index("aspirin"), + "end": text.index("aspirin") + len("aspirin"), + "cui": "387458008", + "value": "aspirin", + }, + { # Incorrect CUI linked + "start": text.index("patient"), + "end": text.index("patient") + len("patient"), + "cui": "116154003", + "value": "patient", + }, + { # Didn't get NER'd + "start": text.index("paracetamol"), + "end": text.index("paracetamol") + len("paracetamol"), + "cui": "387517004", + "value": "paracetamol" + } + ] + return { + "name": "dummy_project", + "id": "0", + "cuis": "", + "tuis": None, + "documents": [{ + "name": "dummy_doc", + "id": "0", + "text": text, + "annotations": annotations, + }], + } + + +fake_cat = DummyCAT() +test_projects = {"projects": [make_fake_test_project()]} + + +class StatsTests(TrainedModelTests): @classmethod def setUpClass(cls): - super().setUpClass() - with open(cls.PERFECT_STATS_PATH) as f: - cls.data: MedCATTrainerExport = json.load(f) - (cls.fps, cls.fns, cls.tps, cls.prec, cls.rec, cls.f1, - cls.counts, cls.examples) = stats.get_stats(cls.model, cls.data) - - def _iter_anns(self): - for proj in self.data["projects"]: - for doc in proj["documents"]: - text = doc["text"] - for ann in doc["annotations"]: - yield proj, doc, text, ann - - # just a sanity check - def test_check_export_is_valid(self): - for proj, doc, text, ann in self._iter_anns(): - start, end, value = ann["start"], ann['end'], ann['value'] - with self.subTest(f"{proj['name']} ({proj['id']}): " - f"{doc['name']} ({doc['id']}) -> " - f"{ann['cui']} ({value}) @ " - f"{start}...{end}"): - detexted_value = text[start:end] - self.assertEqual(detexted_value, value) - - def assert_perfect_dict(self, d: dict[str, Union[float, int]]) -> None: - for cui, f1 in d.items(): - with self.subTest(cui): - self.assertEqual(f1, 1) - - def test_gets_perfect_f1(self): - self.assert_perfect_dict(self.f1) - - def test_gets_perfect_prec(self): - self.assert_perfect_dict(self.prec) - - def test_gets_perfect_rec(self): - self.assert_perfect_dict(self.rec) - - def test_no_fps(self): - self.assertFalse(self.fps) - - def test_no_fns(self): - self.assertFalse(self.fns) - - def test_has_counts_for_concepts(self): - for cui in self.model.cdb.cui2info: - with self.subTest(cui): - cnts = self.counts.get(cui, 0) - self.assertGreater(cnts, 0) + cls.cat = DummyCAT() + cls.cat_linker = DummyCATLinker() + cls.data = {"projects": [make_fake_test_project()]} + cls.result = stats.get_stats_calculator( + cat=cls.cat, + data=cls.data, + use_project_filters=False, + ner_performance=True, + linking_performance=False, + do_print=False, + ) + cls.linker_result = stats.get_stats_calculator( + cat=cls.cat_linker, + data=cls.data, + use_project_filters=False, + ner_performance=False, + linking_performance=True, + do_print=False, + ) + + def test_basic_counts(self) -> None: + stats = self.result.stats.all_projects.get_mode(MetricMode.FULL).stats + # Raw counts of the full pipeline + self.assertEqual(stats.cui_gold_counts["195967001"], 1) + self.assertEqual(stats.cui_gold_counts["387458008"], 1) + self.assertEqual(stats.no_tokens, 0) + self.assertDictEqual(stats.cui_no_tokens, {}) + + ner_stats = self.result.stats.all_projects.get_mode(MetricMode.NER).stats + # Raw counts of the NER only mode + self.assertEqual(ner_stats.cui_gold_counts["__NER__"], 4) + + def test_raw_counts_full_pipe(self) -> None: + stats = self.result.stats.all_projects.get_mode(MetricMode.FULL).stats + # What we got correct + self.assertEqual(stats.tp, 2) + self.assertEqual(stats.cui_tp["195967001"], 1) + self.assertEqual(stats.cui_tp["387458008"], 1) + # The patient error, wrong linked CUI + self.assertEqual(stats.fp, 1) + self.assertEqual(stats.fn, 2) + self.assertEqual(stats.cui_fp["25609006"], 1) + self.assertEqual(stats.cui_fn["116154003"], 1) + + def test_raw_counts_ner_only(self) -> None: + stats = self.result.stats.all_projects.get_mode(MetricMode.NER).stats + # NER only will correctly fix the patient error, as it doesn't care about the CUI, just the span + self.assertEqual(stats.tp, 3) + self.assertEqual(stats.fp, 0) + self.assertEqual(stats.fn, 1) + + def test_raw_counts_linking_only(self) -> None: + stats = self.linker_result.stats.all_projects.get_mode(MetricMode.LINKING).stats + # it's not easily possible to test the linker + # as predictions in the dummy set are hard coded + self.assertEqual(stats.tp, 3) + self.assertEqual(stats.fp, 2) + self.assertEqual(stats.fn, 1) + + def test_precision_recall_f1(self) -> None: + # Full pipeline + metrics = self.result.stats.all_projects.get_mode(MetricMode.FULL).metrics.overall + self.assertAlmostEqual(metrics.precision, 2/3) + self.assertAlmostEqual(metrics.recall, 2/4) + self.assertAlmostEqual(metrics.f1, 0.57, places=2) + + # NER only + ner_pipe = self.result.stats.all_projects.get_mode(MetricMode.NER).metrics.overall + self.assertAlmostEqual(ner_pipe.precision, 3/3) + self.assertAlmostEqual(ner_pipe.recall, 3/4) + self.assertAlmostEqual(ner_pipe.f1, 0.85, places=1) + + # Linking only + linking_pipe = self.linker_result.stats.all_projects.get_mode(MetricMode.LINKING).metrics.overall + self.assertAlmostEqual(linking_pipe.precision, 0.6) + self.assertAlmostEqual(linking_pipe.recall, 3/4) + self.assertAlmostEqual(linking_pipe.f1, 0.666, places=2) + + def test_per_cui_precision_recall_f1(self) -> None: + full_pipe = self.result.stats.all_projects.get_mode(MetricMode.FULL).metrics.per_cui + for cui in ["195967001", "387458008"]: + self.assertAlmostEqual(full_pipe[cui].precision, 1.0) + self.assertAlmostEqual(full_pipe[cui].recall, 1.0) + self.assertAlmostEqual(full_pipe[cui].f1, 1.0) + + for cui in ["25609006", "116154003"]: + self.assertAlmostEqual(full_pipe[cui].precision, 0.0) + self.assertAlmostEqual(full_pipe[cui].recall, 0.0) + self.assertAlmostEqual(full_pipe[cui].f1, 0.0) + + ner_pipe = self.result.stats.all_projects.get_mode(MetricMode.NER).metrics.per_cui + self.assertAlmostEqual(ner_pipe["__NER__"].precision, 1.0) + self.assertAlmostEqual(ner_pipe["__NER__"].recall, 0.75) + self.assertAlmostEqual(ner_pipe["__NER__"].f1, 0.85, places=1) + + def test_cuis_exist(self) -> None: + cui_metrics = self.result.stats.all_projects.get_mode(MetricMode.FULL).metrics.per_cui + ner_cui_metrics = self.result.stats.all_projects.get_mode(MetricMode.NER).metrics.per_cui + linker_cui_metrics = self.linker_result.stats.all_projects.get_mode(MetricMode.LINKING).metrics.per_cui + self.assertIn("__NER__", ner_cui_metrics) + self.assertNotIn("__NER__", cui_metrics) + for cui in ["195967001", "387458008", "25609006", "116154003", "387517004"]: + self.assertNotIn(cui, ner_cui_metrics) + self.assertIn(cui, cui_metrics) + self.assertIn(cui, linker_cui_metrics) + + def test_character_statistics(self) -> None: + full_metrics = self.result.stats.all_projects.get_mode(MetricMode.FULL).metrics.overall + # two cuis are perfect 1 + 1 = 2 + # two are incorrect 2 intersection, 5 union = 0.4 + self.assertAlmostEqual(full_metrics.char_iou, 0.4) + # two are incorrect 2 intersection, 4 union = 0.5 + self.assertAlmostEqual(full_metrics.char_giou, 0.5) + self.assertAlmostEqual(full_metrics.char_cohen_k, 0.45, places=1) + + ner_metrics = self.result.stats.all_projects.get_mode(MetricMode.NER).metrics.overall + # there's only one CUI, so it's the length calculations as below + intersection = len("asthma") + len("aspirin") + len("patient") + union = len("asthma") + len("aspirin") + len("patient") + len("paracetamol") + self.assertAlmostEqual(ner_metrics.char_iou, intersection/union) + self.assertAlmostEqual(ner_metrics.char_giou, intersection/union) + self.assertAlmostEqual(ner_metrics.char_cohen_k, 0.63, places=2) + + linking_metrics = self.linker_result.stats.all_projects.get_mode(MetricMode.LINKING).metrics.overall + self.assertAlmostEqual(linking_metrics.char_iou, 0.6) + # one is incorrect 2 intersection, 4 union = 0.5 + self.assertAlmostEqual(linking_metrics.char_giou, 0.75) + self.assertAlmostEqual(linking_metrics.char_cohen_k, 0.6, places=1) + + def test_per_cui_character_statistics(self) -> None: + full_metrics = self.result.stats.all_projects.get_mode(MetricMode.FULL).metrics.per_cui + # 195967001 and 387458008 are perfect, so IoU = 1 + self.assertAlmostEqual(full_metrics["195967001"].char_iou, 1.0) + self.assertAlmostEqual(full_metrics["387458008"].char_iou, 1.0) + # 25609006 and 116154003 are incorrect, so IoU = 0 + self.assertAlmostEqual(full_metrics["25609006"].char_iou, 0.0) + self.assertAlmostEqual(full_metrics["116154003"].char_iou, 0.0) + + ner_metrics = self.result.stats.all_projects.get_mode(MetricMode.NER).metrics.per_cui + # there's only one CUI, so it's the length calculations as below + # same as previous! + intersection = len("asthma") + len("aspirin") + len("patient") + union = len("asthma") + len("aspirin") + len("patient") + len("paracetamol") + self.assertAlmostEqual(ner_metrics["__NER__"].char_iou, intersection/union) + self.assertAlmostEqual(ner_metrics["__NER__"].char_giou, intersection/union) + self.assertAlmostEqual(ner_metrics["__NER__"].char_cohen_k, 0.63, places=2) + + linker_metrics = self.linker_result.stats.all_projects.get_mode(MetricMode.LINKING).metrics.per_cui + self.assertAlmostEqual(linker_metrics["195967001"].char_iou, 1.0) + self.assertAlmostEqual(linker_metrics["387458008"].char_iou, 1.0) + self.assertAlmostEqual(linker_metrics["25609006"].char_iou, 0.0) + self.assertAlmostEqual(linker_metrics["116154003"].char_iou, 0.0) + self.assertAlmostEqual(linker_metrics["387517004"].char_iou, 1.0) + \ No newline at end of file diff --git a/medcat-v2/tests/utils/test_training_utils.py b/medcat-v2/tests/utils/test_training_utils.py index 3bafac66c..975c5dda2 100644 --- a/medcat-v2/tests/utils/test_training_utils.py +++ b/medcat-v2/tests/utils/test_training_utils.py @@ -23,6 +23,7 @@ def __init__(self, start: int, end: int, text: str, cui: str): self.base = _FakeEntityBase(start, end, text) self.cui = cui self.context_similarity = 1.0 + self.id = 0 class _FakeToken: