anonymous_demo/functional/checkpoint/checkpoint_manager.py

import os
from findfile import find_file

from anonymous_demo.core.tad.prediction.tad_classifier import TADTextClassifier
from anonymous_demo.utils.demo_utils import retry


class CheckpointManager:
    pass


class TADCheckpointManager(CheckpointManager):
    @staticmethod
    @retry
    def get_tad_text_classifier(checkpoint: str = None, eval_batch_size=128, **kwargs):
        tad_text_classifier = TADTextClassifier(
            checkpoint, eval_batch_size=eval_batch_size, **kwargs
        )
        return tad_text_classifier