Upload 9 files

DiT_inference.py

@@ -0,0 +1,99 @@
+import os
+import numpy as np
+from tqdm import tqdm
+import torch
+from datasets import load_dataset, ClassLabel
+from datasets import Features, Array3D
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from metrics import apply_metrics
+
+
+def process_label_ids(batch, remapper, label_column="label"):
+    batch[label_column] = [remapper[label_id] for label_id in batch[label_column]]
+    return batch
+
+
+CACHE_DIR = "/mnt/lerna/data/HFcache" if os.path.exists("/mnt/lerna/data/HFcache") else None
+
+
+def main(args):
+    dataset = load_dataset(args.dataset, split="test", cache_dir=CACHE_DIR)
+    if args.dataset == "rvl_cdip":
+        dataset = dataset.select([i for i in range(len(dataset)) if i != 33669])  # corrupt sample
+    batch_size = 100 if args.dataset == "jordyvl/RVL-CDIP-N" else 1000
+
+    feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
+    model = AutoModelForImageClassification.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
+
+    label2idx = {label.replace(" ", "_"): i for label, i in model.config.label2id.items()}
+    data_idx2label = dict(zip(enumerate(dataset.features["label"].names)))
+    data_label2idx = {label: i for i, label in enumerate(dataset.features["label"].names)}
+    model_idx2label = dict(zip(label2idx.values(), label2idx.keys()))
+    diff = [i for i in range(len(data_label2idx)) if data_idx2label[i] != model_idx2label[i]]
+
+    if diff:
+        print(f"aligning labels {diff}")
+        print(f"model labels: {model_idx2label}")
+        print(f"data labels: {data_idx2label}")
+        print(f"Remapping to {label2idx}")  # YET cannot change the length of the labels
+
+        remapper = {}
+        for k, v in label2idx.items():
+            if k in data_label2idx:
+                remapper[data_label2idx[k]] = v
+
+        print(remapper)
+        new_features = Features(
+            {
+                **{k: v for k, v in dataset.features.items() if k != "label"},
+                "label": ClassLabel(num_classes=len(label2idx), names=list(label2idx.keys())),
+            }
+        )
+
+        dataset = dataset.map(
+            lambda example: process_label_ids(example, remapper),
+            features=new_features,
+            batched=True,
+            batch_size=batch_size,
+            desc="Aligning the labels",
+        )
+
+    features = Features({**dataset.features, "pixel_values": Array3D(dtype="float32", shape=(3, 224, 224))})
+
+    encoded_dataset = dataset.map(
+        lambda examples: feature_extractor([image.convert("RGB") for image in examples["image"]]),
+        batched=True,
+        batch_size=batch_size,
+        features=features,
+    )
+    encoded_dataset.set_format(type="torch", columns=["pixel_values", "label"])
+    BATCH_SIZE = 16
+    dataloader = torch.utils.data.DataLoader(encoded_dataset, batch_size=BATCH_SIZE)
+
+    all_logits, all_references = np.zeros((len(encoded_dataset), len(label2idx))), np.zeros(
+        len(encoded_dataset), dtype=int
+    )
+
+    count = 0
+    for i, batch in tqdm(enumerate(dataloader)):
+        with torch.no_grad():
+            outputs = model(batch["pixel_values"])
+            logits = outputs.logits
+            all_logits[count : count + BATCH_SIZE] = logits.detach().cpu().numpy()
+            all_references[count : count + BATCH_SIZE] = batch["label"].detach().cpu().numpy()
+            count += len(batch["label"])
+
+    all_references = np.array(all_references)
+    all_logits = np.array(all_logits)
+    results = apply_metrics(all_references, all_logits)
+    print(results)
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser("""DiT inference on dataset test set""")
+    parser.add_argument("-d", dest="dataset", type=str, default="rvl_cdip", help="the dataset to be evaluated")
+    args = parser.parse_args()
+
+    main(args)

README.md

@@ -0,0 +1,18 @@
+# Beyond Document Page Classification
+
+## Installation
+
+The scripts require [python >= 3.8](https://www.python.org/downloads/release/python-380/) to run.
+We will create a fresh virtualenvironment in which to install all required packages.
+```sh
+mkvirtualenv -p /usr/bin/python3 BYD
+```
+
+Using poetry and the readily defined pyproject.toml, we will install all required packages
+```sh
+workon BYD 
+pip3 install poetry
+poetry install
+```
+
+## something

experiments.sh

@@ -0,0 +1,23 @@
+# Single page f_p experiments
+command="python3 DiT_inference.py -d jordyvl/RVL_CDIP_N" ; echo $command ; $command
+command="python3 DiT_inference.py -d jordyvl/RVL_CDIP_N" ; echo $command ; $command
+
+
+# Multi-page f_d experiments
+command="python3 simulate_document_classifier.py grid -d bdpc/rvl_cdip_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py first -d bdpc/rvl_cdip_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py second -d bdpc/rvl_cdip_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py last -d bdpc/rvl_cdip_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py max_confidence -d bdpc/rvl_cdip_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py soft_voting -d bdpc/rvl_cdip_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py hard_voting -d bdpc/rvl_cdip_mp" ; echo $command ; $command
+
+
+command="python3 simulate_document_classifier.py grid -d bdpc/rvl_cdip_n_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py first -d bdpc/rvl_cdip_n_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py second -d bdpc/rvl_cdip_n_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py last -d bdpc/rvl_cdip_n_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py hard_voting -d bdpc/rvl_cdip_n_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py soft_voting -d bdpc/rvl_cdip_n_mp" ; echo $command ; $command
+command="python3 simulate_document_classifier.py max_confidence -d bdpc/rvl_cdip_n_mp" ; echo $command ; $command
+

inference_methods.py

@@ -0,0 +1,114 @@
+from enum import Enum
+from PIL import Image
+
+
+def unravel_index(index, shape):
+    out = []
+    for dim in reversed(shape):
+        out.append(index % dim)
+        index = index // dim
+    return tuple(reversed(out))
+
+
+class ExplicitEnum(Enum):
+    """
+    Enum with more explicit error message for missing values or getting all options
+    """
+
+    @classmethod
+    def _missing_(cls, value):
+        raise ValueError(
+            f"{value} is not a valid {cls.__name__}, please select one of {list(cls._value2member_map_.keys())}"
+        )
+
+    @classmethod
+    def options(cls):
+        return list(cls._value2member_map_.keys())
+
+
+class InferenceMethod(ExplicitEnum):
+    """All the implemented inference methods"""
+
+    FIRST = "first"  # check on data setup
+    SECOND = "second"  # robustness for multipage categories
+    LAST = "last"  # robustness for multipage categories
+
+    GRID = "grid"  # create a grid (equal spaced/ OCR density based)
+    # downscale resolution ; might be fine for classification
+
+    MAX_CONFIDENCE = "max_confidence"  # page with highest confidence overall
+    SOFT_VOTING = "soft_voting"  # sum conf/N   -> logits/softmax
+    HARD_VOTING = "hard_voting"  # count votes
+
+    @property
+    def scope(self):
+        if self in [InferenceMethod.FIRST, InferenceMethod.SECOND, InferenceMethod.LAST]:
+            return "sample"
+        if self in [InferenceMethod.GRID]:
+            return "sample-grid"  # single image yet transformation required
+        else:
+            return "iter"
+
+    def get_page_scope(self, pages):
+        if self.scope == "iter":
+            return pages
+        if self == InferenceMethod.GRID:
+            try:
+                return equal_image_grid(pages)
+            except Exception as e:
+                return pages[-1]  # last to not positively bias
+        if self == InferenceMethod.FIRST:
+            return pages[0]
+        if self == InferenceMethod.SECOND:
+            if len(pages) > 1:
+                return pages[1]
+            return pages[0]  # backoff
+        if self == InferenceMethod.LAST:
+            return pages[-1]
+
+    def apply_decision_strategy(self, page_logits):
+        """
+        page logits is of shape [NUM_PAGES x CLASSES]
+        """
+        if self == InferenceMethod.MAX_CONFIDENCE:
+            index = page_logits.argmax()  # tensor with one number
+            indices = unravel_index(index, page_logits.shape)
+            print(f"The page which is max confident: {indices[0]}")
+            return indices[-1]
+        if self == InferenceMethod.HARD_VOTING:
+            return page_logits.argmax(-1).max()
+        if self == InferenceMethod.SOFT_VOTING:
+            return page_logits.mean(0).argmax(-1)
+
+
+def equal_image_grid(images):
+    def compute_grid(n, max_cols=6):
+        equalDivisor = int(n**0.5)
+        cols = min(equalDivisor, max_cols)
+        rows = equalDivisor
+        if rows * cols >= n:
+            return rows, cols
+        cols += 1
+        if rows * cols >= n:
+            return rows, cols
+        while rows * cols < n:
+            rows += 1
+        return rows, cols
+
+    # assert len(images) == rows*cols
+    rows, cols = compute_grid(len(images))
+
+    # rescaling to min width [height padding]
+    images = [im for im in images if (im.height > 0) and (im.width > 0)]  # could be NA
+
+    min_width = min(im.width for im in images)
+    images = [im.resize((min_width, int(im.height * min_width / im.width)), resample=Image.BICUBIC) for im in images]
+
+    w, h = max([img.size[0] for img in images]), max([img.size[1] for img in images])
+
+    grid = Image.new("RGB", size=(cols * w, rows * h))
+    grid_w, grid_h = grid.size
+
+    for i, img in enumerate(images):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid

load_predictions.py

@@ -0,0 +1,180 @@
+import sys
+import numpy as np
+import pandas as pd
+from metrics import ece_logits, aurc_logits, multi_aurc_plot, apply_metrics
+from sklearn.metrics import f1_score
+from collections import OrderedDict
+
+EXPERIMENT_ROOT = "/mnt/lerna/experiments"
+
+
+def softmax(x, axis=-1):
+    # Subtract the maximum value for numerical stability
+    x = x - np.max(x, axis=axis, keepdims=True)
+
+    # Compute the exponentials of the shifted input
+    exps = np.exp(x)
+
+    # Compute the sum of exponentials along the last axis
+    exps_sum = np.sum(exps, axis=axis, keepdims=True)
+
+    # Compute the softmax probabilities
+    softmax_probs = exps / exps_sum
+
+    return softmax_probs
+
+
+def predictions_loader(predictions_path):
+    data = np.load(predictions_path)["arr_0"]
+    dataset_idx = data[:, -1]
+    labels = data[:, -2]
+    if "DiT-base-rvl_cdip_MP" in predictions_path and any(x in predictions_path for x in ["first", "second", "last"]):
+        data = data[:, :-2]  # logits
+        predictions = np.argmax(data, -1)
+    else:
+        labels = data[:, -2].astype(int)
+        predictions = data[:, -3].astype(int)
+        data = data[:, :-3]  # logits
+    return data, labels, predictions, dataset_idx
+
+
+def compare_errors():
+    """
+    from scipy.stats import pearsonr, spearmanr
+    #idx = [x for x in strategy_correctness['first'] if x ==0]
+    spearmanr(strategy_correctness['first'], strategy_correctness['second'])
+    #SignificanceResult(statistic=0.5429413617297623, pvalue=0.0)
+    spearmanr(strategy_correctness['first'], strategy_correctness['last'])
+    #SignificanceResult(statistic=0.5005224326802595, pvalue=0.0)
+
+    pearsonr(strategy_correctness['first'], strategy_correctness['second'])
+    #PearsonRResult(statistic=0.5429413617297617, pvalue=0.0)
+    pearsonr(strategy_correctness['first'], strategy_correctness['last'])
+    #PearsonRResult(statistic=0.5005224326802583, pvalue=0.0)
+    """
+    for dataset in ["rvl_cdip_n_mp"]:  # "DiT-base-rvl_cdip_MP",
+        strategy_logits = {}
+        strategy_correctness = {}
+        for strategy in ["first", "second", "last"]:
+            path = f"{EXPERIMENT_ROOT}/{dataset}/dit-base-finetuned-rvlcdip_{strategy}-0-final.npz"
+
+            strategy_logits[strategy], labels, predictions, dataset_idx = predictions_loader(path)
+            strategy_correctness[strategy] = (predictions == labels).astype(int)
+
+        print("Base accuracy of first: ", np.mean(strategy_correctness["first"]))
+        firstcorrectifsecondcorrect = [
+            x if x == 1 else strategy_correctness["second"][i] for i, x in enumerate(strategy_correctness["first"])
+        ]  # if x ==0]
+        print(f"Accuracy of first when adding knowledge from second page: {np.mean(firstcorrectifsecondcorrect)}")
+        firstcorrectiflastcorrect = [
+            x if x == 1 else strategy_correctness["last"][i] for i, x in enumerate(strategy_correctness["first"])
+        ]  # if x ==0]
+        print(f"Accuracy of first when adding knowledge from last page: {np.mean(firstcorrectiflastcorrect)}")
+
+        firstcorrectifsecondorlastcorrect = [
+            x if x == 1 else (strategy_correctness["second"][i] or strategy_correctness["last"][i])
+            for i, x in enumerate(strategy_correctness["first"])
+        ]  # if x ==0]
+        print(
+            f"Accuracy of first when adding knowledge from second/last page: {np.mean(firstcorrectifsecondorlastcorrect)}"
+        )
+
+        # inverse
+        print("Base accuracy of second: ", np.mean(strategy_correctness["second"]))
+        secondcorrectiffirstcorrect = [
+            x if x == 1 else strategy_correctness["first"][i] for i, x in enumerate(strategy_correctness["second"])
+        ]  # if x ==0]
+        print(f"Accuracy of second when adding knowledge from first page: {np.mean(secondcorrectiffirstcorrect)}")
+        secondcorrectiflastcorrect = [
+            x if x == 1 else strategy_correctness["last"][i] for i, x in enumerate(strategy_correctness["second"])
+        ]  # if x ==0]
+        print(f"Accuracy of second when adding knowledge from last page: {np.mean(secondcorrectiflastcorrect)}")
+
+        # inverse second
+        print("Base accuracy of last: ", np.mean(strategy_correctness["last"]))
+        lastcorrectiffirstcorrect = [
+            x if x == 1 else strategy_correctness["first"][i] for i, x in enumerate(strategy_correctness["last"])
+        ]  # if x ==0]
+        print(f"Accuracy of last when adding knowledge from first page: {np.mean(lastcorrectiffirstcorrect)}")
+        lastcorrectifsecondcorrect = [
+            x if x == 1 else strategy_correctness["second"][i] for i, x in enumerate(strategy_correctness["last"])
+        ]  # if x ==0]
+        print(f"Accuracy of last when adding knowledge from second page: {np.mean(lastcorrectifsecondcorrect)}")
+
+
+def review_one(path):
+    collect = OrderedDict()
+    try:
+        logits, labels, predictions, dataset_idx = predictions_loader(path)
+    except Exception as e:
+        print(f"something went wrong in inference loading {e}")
+        return
+    # print(logits.shape, labels.shape, logits[-1], labels[-1], dataset_idx[-1])
+    y_correct = (predictions == labels).astype(int)
+    acc = np.mean(y_correct)
+    p_hat = np.array([softmax(p, -1)[predictions[i]] for i, p in enumerate(logits)])
+
+    res = aurc_logits(
+        y_correct, p_hat, plot=False, get_cache=True, use_as_is=True
+    )  # DEV: implementation hack to allow for passing I[Y==y_hat] and p_hat instead of logits and label indices
+
+    collect["aurc"] = res["aurc"]
+    collect["accuracy"] = 100 * acc
+    collect["f1"] = 100 * f1_score(labels, predictions, average="weighted")
+    collect["f1_macro"] = 100 * f1_score(labels, predictions, average="macro")
+    collect["ece"] = ece_logits(np.logical_not(y_correct), np.expand_dims(p_hat, -1), use_as_is=True)
+
+    df = pd.DataFrame.from_dict([collect])
+    # df = df[["accuracy", "f1", "f1_macro", "ece", "aurc"]]
+    print(df.to_latex())
+    print(df.to_string())
+    return collect, res
+
+
+def experiments_review():
+    STRATEGIES = ["first", "second", "last", "max_confidence", "soft_voting", "hard_voting", "grid"]
+    for dataset in ["DiT-base-rvl_cdip_MP", "rvl_cdip_n_mp"]:
+        collect = {}
+        aurcs = []
+        caches = []
+        for strategy in STRATEGIES:
+            path = f"{EXPERIMENT_ROOT}/{dataset}/dit-base-finetuned-rvlcdip_{strategy}-0-final.npz"
+            collect[strategy], res = review_one(path)
+            aurcs.append(res["aurc"])
+            caches.append(res["cache"])
+
+        df = pd.DataFrame.from_dict(collect, orient="index")
+        df = df[["accuracy", "f1", "f1_macro", "ece", "aurc"]]
+        print(df.to_latex())
+        print(df.to_string())
+        """
+        subset = [0, 1, 2]
+        multi_aurc_plot(
+            [x for i, x in enumerate(caches) if i in subset],
+            [x for i, x in enumerate(STRATEGIES) if i in subset],
+            aurcs=[x for i, x in enumerate(aurcs) if i in subset],
+        )
+        """
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser("""Deeper evaluation of different inference strategies to classify a document""")
+    DEFAULT = "./dit-base-finetuned-rvlcdip_last-10.npz"
+    parser.add_argument(
+        "predictions_path",
+        type=str,
+        default=DEFAULT,
+        nargs="?",
+        help="path to predictions",
+    )
+
+    args = parser.parse_args()
+    if args.predictions_path == DEFAULT:
+        experiments_review()
+        compare_errors()
+        sys.exit(1)
+
+    print(f"Running default experiment on {args.predictions_path}")
+    review_one(args.predictions_path)

mapping_functions.py

@@ -0,0 +1,157 @@
+import os
+from io import BytesIO
+from tqdm import tqdm
+import numpy as np
+from typing import Callable, Dict, List  # , Literal, NamedTuple, Optional, Tuple, Union
+from PIL import Image as PIL_Image
+from PIL.Image import Image
+
+from datasets import logging
+
+logger = logging.get_logger(__name__)
+import PyPDF2
+
+MAX_PAGES = 50
+MAX_PDF_SIZE = 100000000  # almost 100MB
+MIN_WIDTH, MIN_HEIGHT = 150, 150
+import pdf2image
+
+
+def pdf2image_image_extraction(pdf_stream):
+    try:
+        images: List[Image] = pdf2image.convert_from_bytes(pdf_stream)
+        return images
+    except Exception as e:
+        logger.warning(f"{e}")
+
+
+def pdf_to_pixelvalues_extractor(example, feature_extractor, inference_method):
+    example["pages"] = 0
+    example["pixel_values"] = None
+    pixel_values = []
+    if len(example["file"]) > MAX_PDF_SIZE:
+        logger.warning(f"too large file {len(example['file'])}")
+        return example
+    try:
+        reader = PyPDF2.PdfReader(BytesIO(example["file"]))
+    except Exception as e:
+        logger.warning(f"read_pdf {e}")
+        return example
+    example["pages"] = len(reader.pages)
+    reached_page_limit = False
+    if "sample" in inference_method.scope and inference_method.scope != "sample-grid":
+        page_iterator = [inference_method.get_page_scope(reader.pages)]
+    else:
+        page_iterator = reader.pages
+
+    try:
+        for p, page in enumerate(page_iterator):
+            if reached_page_limit:
+                break
+            for image in page.images:
+                if len(pixel_values) == MAX_PAGES:
+                    reached_page_limit = True
+                    break
+                im = PIL_Image.open(BytesIO(image.data))
+                if im.width < MIN_WIDTH and im.height < MIN_HEIGHT:
+                    continue
+                # try:
+                # except Exception as e:
+                #     logger.warning(f"get_images {e}")
+                if inference_method.scope != "sample-grid":
+                    im = feature_extractor([im.convert("RGB")])["pixel_values"][0]
+                pixel_values.append(im)
+    except Exception as e:
+        print(f"{example.get('id')} PyPDF get_images {e}")
+        pixel_values = []
+
+    if len(pixel_values) == 0:
+        # at least try with another API
+        try:
+            images = pdf2image_image_extraction(example["file"])
+        except Exception as e:
+            print(f"{example.get('id')} pdf2image get_images {e}")
+            images = []
+
+        if not images:
+            print(f"{example.get('id')} pdf2image has no images")
+            example["pages"] = 0
+            return example
+
+        # got lucky with pdf2image
+        example["pages"] = len(images)
+        for im in images:
+            if len(pixel_values) == MAX_PAGES:
+                reached_page_limit = True
+                break
+            if im.width < MIN_WIDTH and im.height < MIN_HEIGHT:
+                continue
+            if inference_method.scope != "sample-grid":
+                im = feature_extractor([im.convert("RGB")])["pixel_values"][0]
+            pixel_values.append(im)
+
+    if inference_method.scope == "sample-grid":
+        grid = inference_method.get_page_scope(pixel_values)
+        pixel_values = feature_extractor([grid.convert("RGB")])["pixel_values"][0]
+    elif "sample" in inference_method.scope:
+        pixel_values = pixel_values[0]
+    example["pixel_values"] = np.array(pixel_values)
+    return example
+
+
+def nativepdf_to_pixelvalues_extractor(example, feature_extractor, inference_method):
+    IMPOSSIBLE = ["6483941-Letter-to-John-Campbell.pdf", "7276809-Ocoee-Newspaper-Pages.pdf"]
+    example["pages"] = 0
+    example["pixel_values"] = None
+    pixel_values = []
+    if len(example["file"]) > MAX_PDF_SIZE:
+        logger.warning(f"too large file {len(example['file'])}")
+        return example
+
+    # images = example['images']
+    try:
+        images = pdf2image_image_extraction(example["file"])
+    except Exception as e:
+        print(f"{example.get('id')} pdf2image get_images {e}")
+        images = []
+
+    if not images:
+        print(f"{example.get('id')} pdf2image has no images")
+        example["pages"] = 0
+        return example
+
+    # do image checks before and after
+    images = [im for im in images if im.width >= MIN_WIDTH and im.height >= MIN_HEIGHT]
+
+    if not images or (example.get("id") in IMPOSSIBLE and inference_method.scope == "sample-grid"):
+        print(f"{example.get('id')} pdf2image has no images")
+        example["pages"] = 0
+        return example
+
+    example["pages"] = len(images)
+    reached_page_limit = False
+    if "sample" in inference_method.scope and inference_method.scope != "sample-grid":
+        page_iterator = [inference_method.get_page_scope(images)]
+    else:
+        page_iterator = images
+
+    for im in page_iterator:
+        if len(pixel_values) == MAX_PAGES:
+            reached_page_limit = True
+            break
+        if inference_method.scope != "sample-grid":
+            im = feature_extractor([im.convert("RGB")])["pixel_values"][0]
+        pixel_values.append(im)
+
+    if len(pixel_values) == 0:
+        print(f"{example.get('id')} pdf2image has no valid images")
+        example["pages"] = 0
+        return example
+
+    if inference_method.scope == "sample-grid":
+        grid = inference_method.get_page_scope(pixel_values)
+        pixel_values = feature_extractor([grid.convert("RGB")])["pixel_values"][0]
+    elif "sample" in inference_method.scope:
+        pixel_values = pixel_values[0]
+    example["pixel_values"] = np.array(pixel_values)
+    return example

metrics.py

@@ -0,0 +1,309 @@
+# base calibration metric
+
+# https://github.com/JonathanWenger/pycalib/blob/master/pycalib/scoring.py
+# https://github.com/google-research/robustness_metrics/blob/master/robustness_metrics/metrics/uncertainty.py
+#  https://github.com/kjdhfg/fd-shifts
+
+from __future__ import annotations
+import scipy
+import sklearn.utils.validation
+
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Any
+import pandas as pd
+import numpy as np
+import numpy.typing as npt
+from collections import OrderedDict
+from sklearn import metrics as skm
+import evaluate as HF_evaluate
+
+ArrayType = npt.NDArray[np.floating]
+
+
+## https://github.com/IML-DKFZ/fd-shifts/blob/main/fd_shifts/analysis/confid_scores.py#L20
+
+# ----------------------------------- general metrics with consistent metric(y_true, p_hat) API -----------------------------------
+
+
+def f1_w(y_true, p_hat, y_hat=None):
+    if y_hat is None:
+        y_hat = np.argmax(p_hat, axis=-1)
+    return skm.f1_score(y_true, y_hat, average="weighted")
+
+
+def f1_micro(y_true, p_hat, y_hat=None):
+    if y_hat is None:
+        y_hat = np.argmax(p_hat, axis=-1)
+    return skm.f1_score(y_true, y_hat, average="micro")
+
+
+def f1_macro(y_true, p_hat, y_hat=None):
+    if y_hat is None:
+        y_hat = np.argmax(p_hat, axis=-1)
+    return skm.f1_score(y_true, y_hat, average="macro")
+
+
+# Pure numpy and TF implementations of proper losses (as metrics) -----------------------------------
+
+
+def brier_loss(y_true, p_hat):
+    r"""Brier score.
+    If the true label is k, while the predicted vector of probabilities is
+    [y_1, ..., y_n], then the Brier score is equal to
+    \sum_{i != k} y_i^2 + (y_k - 1)^2.
+
+    The smaller the Brier score, the better, hence the naming with "loss".
+    Across all items in a set N predictions, the Brier score measures the
+    mean squared difference between (1) the predicted probability assigned
+    to the possible outcomes for item i, and (2) the actual outcome.
+    Therefore, the lower the Brier score is for a set of predictions, the
+    better the predictions are calibrated. Note that the Brier score always
+    takes on a value between zero and one, since this is the largest
+    possible difference between a predicted probability (which must be
+    between zero and one) and the actual outcome (which can take on values
+    of only 0 and 1). The Brier loss is composed of refinement loss and
+    calibration loss.
+
+    """
+    N = len(y_true)
+    K = p_hat.shape[-1]
+
+    if y_true.shape != p_hat.shape:
+        zeros = scipy.sparse.lil_matrix((N, K))
+        for i in range(N):
+            zeros[i, y_true[i]] = 1
+
+    if not np.isclose(np.sum(p_hat), len(p_hat)):
+        p_hat = scipy.special.softmax(p_hat, axis=-1)
+
+    return np.mean(np.sum(np.array(p_hat - zeros) ** 2, axis=1))
+
+
+def nll(y_true, p_hat):
+    r"""Multi-class negative log likelihood.
+    If the true label is k, while the predicted vector of probabilities is
+    [p_1, ..., p_K], then the negative log likelihood is -log(p_k).
+    Does not require onehot encoding
+    """
+    labels = np.arange(p_hat.shape[-1])
+    return skm.log_loss(y_true, p_hat, labels=labels)
+
+
+def accuracy(y_true, p_hat):
+    y_pred = np.argmax(p_hat, axis=-1)
+    return sklearn.metrics.accuracy_score(y_true=y_true, y_pred=y_pred)
+
+
+AURC_DISPLAY_SCALE = 1  # 1000
+
+"""
+From: https://web.stanford.edu/class/archive/cs/cs224n/cs224n.1204/reports/custom/report52.pdf 
+
+The risk-coverage (RC) curve [28, 16] is a measure of the trade-off between the
+coverage (the proportion of test data encountered), and the risk (the error rate under this coverage). Since each
+prediction comes with a confidence score, given a list of prediction correctness Z paired up with the confidence
+scores C, we sort C in reverse order to obtain sorted C'
+, and its corresponding correctness Z'
+. Note that the correctness is computed based on Exact Match (EM) as described in [22]. The RC curve is then obtained by
+computing the risk of the coverage from the beginning of Z'
+(most confident) to the end (least confident). In particular, these metrics evaluate 
+the relative order of the confidence score, which means that we want wrong
+answers have lower confidence score than the correct ones, ignoring their absolute values. 
+
+Source: https://github.com/kjdhfg/fd-shifts 
+
+References:
+-----------
+
+[1] Jaeger, P.F., Lüth, C.T., Klein, L. and Bungert, T.J., 2022. A Call to Reflect on Evaluation Practices for Failure Detection in Image Classification. arXiv preprint arXiv:2211.15259.
+
+[2] Kamath, A., Jia, R. and Liang, P., 2020. Selective Question Answering under Domain Shift. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics (pp. 5684-5696).
+
+"""
+
+
+@dataclass
+class StatsCache:
+    """Cache for stats computed by scikit used by multiple metrics.
+
+    Attributes:
+        confids (array_like): Confidence values
+        correct (array_like): Boolean array (best converted to int) where predictions were correct
+    """
+
+    confids: npt.NDArray[Any]
+    correct: npt.NDArray[Any]
+
+    @cached_property
+    def roc_curve_stats(self) -> tuple[npt.NDArray[Any], npt.NDArray[Any]]:
+        fpr, tpr, _ = skm.roc_curve(self.correct, self.confids)
+        return fpr, tpr
+
+    @property
+    def residuals(self) -> npt.NDArray[Any]:
+        return 1 - self.correct
+
+    @cached_property
+    def rc_curve_stats(self) -> tuple[list[float], list[float], list[float]]:
+        coverages = []
+        risks = []
+
+        n_residuals = len(self.residuals)
+        idx_sorted = np.argsort(self.confids)
+
+        coverage = n_residuals
+        error_sum = sum(self.residuals[idx_sorted])
+
+        coverages.append(coverage / n_residuals)
+        risks.append(error_sum / n_residuals)
+
+        weights = []
+
+        tmp_weight = 0
+        for i in range(0, len(idx_sorted) - 1):
+            coverage = coverage - 1
+            error_sum = error_sum - self.residuals[idx_sorted[i]]
+            selective_risk = error_sum / (n_residuals - 1 - i)
+            tmp_weight += 1
+            if i == 0 or self.confids[idx_sorted[i]] != self.confids[idx_sorted[i - 1]]:
+                coverages.append(coverage / n_residuals)
+                risks.append(selective_risk)
+                weights.append(tmp_weight / n_residuals)
+                tmp_weight = 0
+
+        # add a well-defined final point to the RC-curve.
+        if tmp_weight > 0:
+            coverages.append(0)
+            risks.append(risks[-1])
+            weights.append(tmp_weight / n_residuals)
+        return coverages, risks, weights
+
+
+def aurc(stats_cache: StatsCache):
+    """auc metric function
+    Args:
+        stats_cache (StatsCache): StatsCache object
+    Returns:
+        metric value
+    Important for assessment: LOWER is better!
+    """
+    _, risks, weights = stats_cache.rc_curve_stats
+    return sum([(risks[i] + risks[i + 1]) * 0.5 * weights[i] for i in range(len(weights))]) * AURC_DISPLAY_SCALE
+
+
+def aurc_logits(references, predictions, plot=False, get_cache=False, use_as_is=False):
+    if not use_as_is:
+        if not np.isclose(np.sum(references), len(references)):
+            references = (np.argmax(predictions, -1) == references).astype(int)  # correctness
+
+        if not np.isclose(np.sum(predictions), len(predictions)):
+            predictions = scipy.special.softmax(predictions, axis=-1)
+
+        if predictions.ndim == 2:
+            predictions = np.max(predictions, -1)
+
+    cache = StatsCache(confids=predictions, correct=references)
+
+    if plot:
+        coverages, risks, weights = cache.rc_curve_stats
+        pd.options.plotting.backend = "plotly"
+        df = pd.DataFrame(zip(coverages, risks, weights), columns=["% Coverage", "% Risk", "weights"])
+        fig = df.plot(x="% Coverage", y="% Risk")
+        fig.show()
+    if get_cache:
+        return {"aurc": aurc(cache), "cache": cache}
+    return aurc(cache)
+
+
+def multi_aurc_plot(caches, names, aurcs=None, verbose=False):
+    pd.options.plotting.backend = "plotly"
+    df = pd.DataFrame()
+    for cache, name in zip(caches, names):
+        coverages, risks, weights = cache.rc_curve_stats
+        df[name] = pd.Series(risks, index=coverages)
+    if verbose:
+        print(df.head(), df.index, df.columns)
+    fig = df.plot()
+    title = ""
+    if aurcs is not None:
+        title = "AURC: " + " - ".join([str(round(aurc, 4)) for aurc in aurcs])
+    fig.update_layout(title=title, xaxis_title="% Coverage", yaxis_title="% Risk")
+    fig.show()
+
+
+def ece_logits(references, predictions, use_as_is=False):
+    if not use_as_is:
+        if not np.isclose(np.sum(predictions), len(predictions)):
+            predictions = scipy.special.softmax(predictions, axis=-1)
+
+    metric = HF_evaluate.load("jordyvl/ece")
+    kwargs = dict(
+        n_bins=min(len(predictions) - 1, 100),
+        scheme="equal-mass",
+        bin_range=[0, 1],
+        proxy="upper-edge",
+        p=1,
+        detail=False,
+    )
+
+    ece_result = metric.compute(
+        references=references,
+        predictions=predictions,
+        **kwargs,
+    )
+    return ece_result["ECE"]
+
+
+METRICS = [accuracy, brier_loss, nll, f1_w, f1_macro, ece_logits, aurc_logits]
+
+
+def apply_metrics(y_true, y_probs, metrics=METRICS):
+    predictive_performance = OrderedDict()
+    for metric in metrics:
+        try:
+            predictive_performance[f"{metric.__name__.replace('_logits', '')}"] = metric(y_true, y_probs)
+        except Exception as e:
+            print(e)
+    # print(json.dumps(predictive_performance, indent=4))
+    return predictive_performance
+
+
+def evaluate_coverages(
+    logits, labels, confidence, coverages=[100, 99, 98, 97, 95, 90, 85, 80, 75, 70, 60, 50, 40, 30, 20, 10]
+):
+    correctness = np.equal(logits.argmax(-1), labels)
+    abstention_results = list(zip(list(confidence), list(correctness)))
+    # sort the abstention results according to their reservations, from high to low
+    abstention_results.sort(key=lambda x: x[0])
+    # get the "correct or not" list for the sorted results
+    sorted_correct = list(map(lambda x: int(x[1]), abstention_results))
+    size = len(sorted_correct)
+    print("Abstention Logit: accuracy of coverage ")  # 1-risk
+    for coverage in coverages:
+        covered_correct = sorted_correct[: round(size / 100 * coverage)]
+        print("{:.0f}: {:.3f}, ".format(coverage, sum(covered_correct) / len(covered_correct) * 100.0), end="")
+    print("")
+
+    sr_results = list(zip(list(logits.max(-1)), list(correctness)))
+    # sort the abstention results according to Softmax Response scores, from high to low
+    sr_results.sort(key=lambda x: -x[0])
+    # get the "correct or not" list for the sorted results
+    sorted_correct = list(map(lambda x: int(x[1]), sr_results))
+    size = len(sorted_correct)
+    print("Softmax Response: accuracy of coverage ")
+    for coverage in coverages:
+        covered_correct = sorted_correct[: round(size / 100 * coverage)]
+        print("{:.0f}: {:.3f}, ".format(coverage, sum(covered_correct) / len(covered_correct) * 100.0), end="")
+    print("")
+
+
+def compute_metrics(eval_preds):
+    logits, labels = eval_preds  # output of forward
+    if isinstance(logits, tuple):
+        confidence = logits[1]
+        logits = logits[0]
+        if confidence.size == logits.shape[0]:
+            evaluate_coverages(logits, labels, confidence)
+    results = apply_metrics(labels, logits)
+    return results

pyproject.toml

@@ -0,0 +1,36 @@
+[tool.poetry]
+name = "BDPC"
+version = "0.1.0"
+description = "something"
+authors = ["ANON"]
+
+readme = 'README.md'  # Markdown files are supported
+keywords = ['research', 'evaluation', 'classification', 'documetns']
+
+classifiers = [
+    "Development Status :: 1 - Planning",
+    "Environment :: Web Environment",
+    "Operating System :: POSIX :: Linux",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence"
+]
+
+[tool.poetry.dependencies]
+python = "^3.8"
+numpy = "^1.24.1"
+evaluate = "^0.4.0"
+scikit-learn = "^1.2.0"
+transformers = "^4.26.0"
+pandas = "^1.5.3"
+torch = "^1.13.1"
+pillow = "^9.4.0"
+
+[tool.poetry.dev-dependencies]
+pytest = "^3.10.1"
+pytest-cov = "^2.9.0"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.masonry.api"

simulate_document_classifier.py

@@ -0,0 +1,222 @@
+import os
+from tqdm import tqdm
+import pandas as pd
+import numpy as np
+import torch
+from datasets import load_dataset, logging
+from datasets import Features, Value, Image, Sequence, Array3D, Array4D
+import evaluate
+from metrics import apply_metrics
+
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification  # DiT
+
+logger = logging.get_logger(__name__)
+
+from mapping_functions import (
+    pdf_to_pixelvalues_extractor,
+    nativepdf_to_pixelvalues_extractor,
+)
+from inference_methods import InferenceMethod
+
+EXPERIMENT_ROOT = "/mnt/lerna/experiments"
+
+
+def load_base_model():
+    feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
+    model = AutoModelForImageClassification.from_pretrained("microsoft/dit-base-finetuned-rvlcdip")
+    return model, feature_extractor
+
+
+def logits_monitor(args, running_logits, references, predictions, identifier="a"):
+    output_path = f"{EXPERIMENT_ROOT}/{args.model.split('/')[-1]}_{args.dataset.split('/')[-1]}_{args.inference_method}-{args.downsampling}-i{identifier}.npz"
+
+    raw_output = torch.cat(
+        [
+            torch.cat(running_logits, dim=0).cpu(),
+            torch.Tensor(references).unsqueeze(1),
+            torch.Tensor(predictions).unsqueeze(1),
+            torch.Tensor(np.arange(int(identifier) - len(references), int(identifier))).unsqueeze(1),
+        ],
+        dim=1,
+    )
+    np.savez_compressed(output_path, raw_output.cpu().data.numpy())
+    tqdm.write("saved raw test outputs to {}".format(output_path))
+
+
+def monitor_cleanup(args, buffer_keys):
+    """
+    This merges all previous buffers to 1 file
+    """
+    output_path = f"{EXPERIMENT_ROOT}/{args.model.split('/')[-1]}_{args.dataset.split('/')[-1]}_{args.inference_method}-{args.downsampling}"
+
+    for i, identifier in enumerate(buffer_keys):
+        identifier_path = f"{output_path}-i{identifier}.npz"
+        saved = np.load(identifier_path)["arr_0"]
+        if i == 0:
+            catted = saved
+        else:
+            catted = np.concatenate([catted, saved])
+    out_path = f"{output_path}-final.npz"
+    np.savez_compressed(out_path, catted)
+    tqdm.write("saved raw test outputs to {}".format(out_path))
+    # cleanup
+    for i, identifier in enumerate(buffer_keys):
+        identifier_path = f"{output_path}-i{identifier}.npz"
+        os.remove(identifier_path)
+
+
+def main(args):
+    testds = load_dataset(
+        args.dataset,
+        cache_dir="/mnt/lerna/data/HFcache",
+        split="test",
+        revision=None if args.dataset != "bdpc/rvl_cdip_mp" else "d3a654c9f63f14d0aaa94e08aa30aa3dc20713c1",
+    )
+
+    if args.downsampling:
+        testds = testds.select(list(range(0, args.downsampling)))
+
+    model = AutoModelForImageClassification.from_pretrained(args.model)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+    label2idx = {label: i for label, i in model.config.label2id.items()}  # .replace(" ", "_")
+    print(label2idx)
+
+    data_idx2label = dict(enumerate(testds.features["labels"].names))
+    model_idx2label = dict(zip(label2idx.values(), label2idx.keys()))
+    diff = [i for i in range(len(data_idx2label)) if data_idx2label[i] != model_idx2label[i]]
+    if diff:
+        print(f"aligning labels {diff}")
+        testds = testds.align_labels_with_mapping(label2idx, "labels")
+
+    inference_method = InferenceMethod[args.inference_method.upper()]
+    dummy_inference_method = inference_method
+    feature_extractor = AutoFeatureExtractor.from_pretrained(args.model)
+
+    features = {
+        **{k: v for k, v in testds.features.items() if k in ["labels", "pixel_values", "id"]},
+        "pages": Value(dtype="int32"),
+        "pixel_values": Array3D(dtype="float32", shape=(3, 224, 224)),
+    }
+    if not "sample" in inference_method.scope:
+        features["pixel_values"] = Array4D(dtype="float32", shape=(None, 3, 224, 224))
+        dummy_inference_method = InferenceMethod["max_confidence".upper()]
+    features = Features(features)
+
+    remove_columns = ["file"]
+    if args.dataset == "bdpc/rvl_cdip_mp":
+        image_preprocessor = lambda batch: pdf_to_pixelvalues_extractor(
+            batch, feature_extractor, dummy_inference_method
+        )
+        encoded_testds = testds.map(
+            image_preprocessor, features=features, remove_columns=remove_columns, desc="pdf_to_pixelvalues"
+        )
+    else:
+        image_preprocessor = lambda batch: nativepdf_to_pixelvalues_extractor(
+            batch, feature_extractor, dummy_inference_method
+        )
+        encoded_testds = testds.map(
+            image_preprocessor,
+            features=features,
+            remove_columns=remove_columns,
+            desc="pdf_to_pixelvalues",
+            batch_size=10,
+        )
+        # remove_columns.append("images")
+
+    # select approach
+    print(f"Before filtering: {len(encoded_testds)}")
+    more_complex_filter = lambda example: example["pages"] != 0 and not np.any(np.isnan(example["pixel_values"]))
+    good_indices = [i for i, x in tqdm(enumerate(encoded_testds), desc="filter") if more_complex_filter(x)]
+    encoded_testds = encoded_testds.select(good_indices)
+    print(f"After filtering: {len(encoded_testds)}")
+
+    metric = evaluate.load("accuracy")
+
+    # going to have to manually iterate without dataloader and do tensor conversion
+    encoded_testds.set_format(type="torch", columns=["pixel_values", "labels"])
+    args.batch_size = args.batch_size if "sample" in inference_method.scope else 1
+    dataloader = torch.utils.data.DataLoader(encoded_testds, batch_size=args.batch_size)
+
+    running_logits = []
+    predictions, references = [], []
+    buffer_references = []
+    buffer_predictions = []
+    buffer = 0
+    BUFFER_SIZE = 5000
+    buffer_keys = []
+    for i, batch in tqdm(enumerate(dataloader), desc="Inference loop"):
+        with torch.no_grad():
+            batch["labels"] = batch["labels"].to(device)
+            batch["pixel_values"] = batch["pixel_values"].to(device)
+            if "sample" in inference_method.scope:
+                outputs = model(batch["pixel_values"].to(device))
+                logits = outputs.logits
+                buffer_predictions.extend(logits.argmax(-1).tolist())
+                buffer_references.extend(batch["labels"].tolist())
+                running_logits.append(logits)
+            else:
+                try:
+                    page_logits = model(batch["pixel_values"][0]).logits
+                except Exception as e:
+                    print(f"something went wrong in inference {e}")
+                    continue
+                prediction = inference_method.apply_decision_strategy(page_logits)  # apply logic depending on method
+                buffer_predictions.append(prediction.tolist())
+                buffer_references.extend(batch["labels"].tolist())
+                running_logits.append(page_logits.mean(0).unsqueeze(0))  # average over pages as representative
+
+        buffer += args.batch_size
+        if buffer >= BUFFER_SIZE:
+            predictions.extend(buffer_predictions)
+            references.extend(buffer_references)
+            logits_monitor(args, running_logits, buffer_references, buffer_predictions, identifier=str(i))
+            buffer_keys.append(str(i))
+            running_logits = []
+            buffer_references = []
+            buffer_predictions = []
+            buffer = 0
+
+    if buffer != 0:  # dump remaining out of buffer
+        predictions.extend(buffer_predictions)
+        references.extend(buffer_references)
+        logits_monitor(args, running_logits, buffer_references, buffer_predictions, identifier=str(i))
+        buffer_keys.append(str(i))
+
+    accuracy = metric.compute(references=references, predictions=predictions)
+    print(f"Accuracy on this inference configuration {inference_method}:", accuracy)
+    monitor_cleanup(args, buffer_keys)
+
+
+if __name__ == "__main__":
+    from argparse import ArgumentParser
+
+    parser = ArgumentParser("""Test different inference strategies to classify a document""")
+    parser.add_argument(
+        "inference_method",
+        type=str,
+        default="first",
+        nargs="?",
+        help="how to evaluate DiT on RVL-CDIP_multi",
+    )
+    parser.add_argument("-s", dest="downsampling", type=int, default=0, help="number of testset samples")
+    parser.add_argument("-d", dest="dataset", type=str, default="bdpc/rvl_cdip_mp", help="the dataset to be evaluated")
+    parser.add_argument(
+        "-m",
+        dest="model",
+        type=str,
+        default="microsoft/dit-base-finetuned-rvlcdip",
+        help="the model checkpoint to be evaluated",
+    )
+    parser.add_argument("-b", dest="batch_size", type=int, default=16, help="batch size")
+    parser.add_argument(
+        "-k",
+        dest="keep_in_memory",
+        default=False,
+        action="store_true",
+        help="do not cache operations (for testing)",
+    )
+
+    args = parser.parse_args()
+
+    main(args)