# Buggy code
def load(
    cls,
    model_name_or_path,
    batch_size=4,
    gpu=False,
    task_type=None,
    return_class_probs=False,
    strict=True,
    max_seq_len=256,
    doc_stride=128,
    extraction_layer=None,
    extraction_strategy=None,
):
    """
    Load an Inferencer incl. all relevant components (model, tokenizer, processor ...) either by

    1. specifying a public name from transformers' model hub (https://huggingface.co/models)
    2. or pointing to a local directory it is saved in.

    :param model_name_or_path: Local directory or public name of the model to load.
    :type model_name_or_path: str
    :param batch_size: Number of samples computed once per batch
    :type batch_size: int
    :param gpu: If GPU shall be used
    :type gpu: bool
    :param task_type: Type of task the model should be used for. Currently supporting:
                      "embeddings", "question_answering", "text_classification", "ner". More coming soon...
    :param task_type: str
    :param strict: whether to strictly enforce that the keys loaded from saved model match the ones in
                   the PredictionHead (see torch.nn.module.load_state_dict()).
                   Set to `False` for backwards compatibility with PHs saved with older version of FARM.
    :type strict: bool
    :param max_seq_len: maximum length of one text sample
    :type max_seq_len: int
    :param doc_stride: Only QA: When input text is longer than max_seq_len it gets split into parts, strided by doc_stride
    :type doc_stride: int
    :param extraction_strategy: Strategy to extract vectors. Choices: 'cls_token' (sentence vector), 'reduce_mean'
                           (sentence vector), reduce_max (sentence vector), 'per_token' (individual token vectors)
    :type extraction_strategy: str
    :param extraction_layer: number of layer from which the embeddings shall be extracted. Default: -1 (very last layer).
    :type extraction_layer: int
    :return: An instance of the Inferencer.

    """

    device, n_gpu = initialize_device_settings(
        use_cuda=gpu, local_rank=-1, use_amp=None
    )
    name = os.path.basename(model_name_or_path)

    # a) either from local dir
    if os.path.exists(model_name_or_path):
        model = BaseAdaptiveModel.load(
            load_dir=model_name_or_path, device=device, strict=strict
        )
        if task_type == "embeddings":
            processor = InferenceProcessor.load_from_dir(model_name_or_path)
        else:
            processor = Processor.load_from_dir(model_name_or_path)

    # b) or from remote transformers model hub
    else:
        logger.info(
            f"Could not find `{model_name_or_path}` locally. Try to download from model hub ..."
        )
        if not task_type:
            raise ValueError(
                "Please specify the 'task_type' of the model you want to load from transformers. "
                "Valid options for arg `task_type`:"
                "'question_answering', 'embeddings', 'text_classification', 'ner'"
            )

        model = AdaptiveModel.convert_from_transformers(
            model_name_or_path, device, task_type
        )
        config = AutoConfig.from_pretrained(model_name_or_path)
        tokenizer = Tokenizer.load(model_name_or_path)

        # TODO infer task_type automatically from config (if possible)
        if task_type == "question_answering":
            processor = SquadProcessor(
                tokenizer=tokenizer,
                max_seq_len=max_seq_len,
                label_list=["start_token", "end_token"],
                metric="squad",
                data_dir=None,
                doc_stride=doc_stride,
            )
        elif task_type == "embeddings":
            processor = InferenceProcessor(tokenizer=tokenizer, max_seq_len=max_seq_len)

        elif task_type == "text_classification":
            label_list = list(config.id2label[id] for id in range(len(config.id2label)))
            processor = TextClassificationProcessor(
                tokenizer=tokenizer,
                max_seq_len=max_seq_len,
                data_dir=None,
                label_list=label_list,
                label_column_name="label",
                metric="acc",
                quote_char='"',
            )
        elif task_type == "ner":
            label_list = list(config.label2id.keys())
            processor = NERProcessor(
                tokenizer=tokenizer,
                max_seq_len=max_seq_len,
                data_dir=None,
                metric="seq_f1",
                label_list=label_list,
            )
        else:
            raise ValueError(
                f"`task_type` {task_type} is not supported yet. "
                f"Valid options for arg `task_type`: 'question_answering', "
                f"'embeddings', 'text_classification', 'ner'"
            )

    return cls(
        model,
        processor,
        task_type=task_type,
        batch_size=batch_size,
        gpu=gpu,
        name=name,
        return_class_probs=return_class_probs,
        extraction_strategy=extraction_strategy,
        extraction_layer=extraction_layer,
    )


# Traceback (TypeError)
03/28/2020 22:25:07 - INFO - farm.utils -   device: cpu n_gpu: 0, distributed training: False, automatic mixed precision training: None
03/28/2020 22:25:07 - INFO - farm.modeling.adaptive_model -   Found files for loading 1 prediction heads
03/28/2020 22:25:07 - WARNING - farm.modeling.prediction_head -   Some unused parameters are passed to the QuestionAnsweringHead. Might not be a problem. Params: {"training": true, "num_labels": 2, "ph_output_type": "per_token_squad", "model_type": "span_classification", "name": "QuestionAnsweringHead"}
03/28/2020 22:25:07 - INFO - farm.modeling.prediction_head -   Prediction head initialized with size [768, 2]
03/28/2020 22:25:07 - INFO - farm.modeling.prediction_head -   Loading prediction head from ../../saved_models/twmkn9/albert-base-v2-squad2/prediction_head_0.bin
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/Documents/CodingProjects/NLPofTimFerrissShow/QnA_with_Tim_Haystack.py in
51
52 # Load model
----> 53 reader = FARMReader(model_name_or_path="../../saved_models/twmkn9/albert-base-v2-squad2", use_gpu=False)
54 # A retriever identifies the k most promising chunks of text that might contain the answer for our question
55 # Retrievers use some simple but fast algorithm, here: TF-IDF

/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/haystack/reader/farm.py in __init__(self, model_name_or_path, context_window_size, batch_size, use_gpu, no_ans_boost, top_k_per_candidate, top_k_per_sample, max_processes, max_seq_len, doc_stride)
79         self.inferencer = Inferencer.load(model_name_or_path, batch_size=batch_size, gpu=use_gpu,
80                                           task_type="question_answering", max_seq_len=max_seq_len,
---> 81                                           doc_stride=doc_stride)
82         self.inferencer.model.prediction_heads[0].context_window_size = context_window_size
83         self.inferencer.model.prediction_heads[0].no_ans_boost = no_ans_boost

/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/farm/infer.py in load(cls, model_name_or_path, batch_size, gpu, task_type, return_class_probs, strict, max_seq_len, doc_stride)
139                 processor = InferenceProcessor.load_from_dir(model_name_or_path)
140             else:
--> 141                 processor = Processor.load_from_dir(model_name_or_path)
142
143         # b) or from remote transformers model hub

/Library/Frameworks/Python.framework/Versions/3.7/lib/python3.7/site-packages/farm/data_handler/processor.py in load_from_dir(cls, load_dir)
189         del config["tokenizer"]
190
--> 191         processor = cls.load(tokenizer=tokenizer, processor_name=config["processor"], **config)
192
193         for task_name, task in config["tasks"].items():

TypeError: load() missing 1 required positional argument: 'data_dir'

# Fixed code
def load(
    cls,
    model_name_or_path,
    batch_size=4,
    gpu=False,
    task_type=None,
    return_class_probs=False,
    strict=True,
    max_seq_len=256,
    doc_stride=128,
    extraction_layer=None,
    extraction_strategy=None,
):
    """
    Load an Inferencer incl. all relevant components (model, tokenizer, processor ...) either by

    1. specifying a public name from transformers' model hub (https://huggingface.co/models)
    2. or pointing to a local directory it is saved in.

    :param model_name_or_path: Local directory or public name of the model to load.
    :type model_name_or_path: str
    :param batch_size: Number of samples computed once per batch
    :type batch_size: int
    :param gpu: If GPU shall be used
    :type gpu: bool
    :param task_type: Type of task the model should be used for. Currently supporting:
                      "embeddings", "question_answering", "text_classification", "ner". More coming soon...
    :param task_type: str
    :param strict: whether to strictly enforce that the keys loaded from saved model match the ones in
                   the PredictionHead (see torch.nn.module.load_state_dict()).
                   Set to `False` for backwards compatibility with PHs saved with older version of FARM.
    :type strict: bool
    :param max_seq_len: maximum length of one text sample
    :type max_seq_len: int
    :param doc_stride: Only QA: When input text is longer than max_seq_len it gets split into parts, strided by doc_stride
    :type doc_stride: int
    :param extraction_strategy: Strategy to extract vectors. Choices: 'cls_token' (sentence vector), 'reduce_mean'
                           (sentence vector), reduce_max (sentence vector), 'per_token' (individual token vectors)
    :type extraction_strategy: str
    :param extraction_layer: number of layer from which the embeddings shall be extracted. Default: -1 (very last layer).
    :type extraction_layer: int
    :return: An instance of the Inferencer.

    """

    device, n_gpu = initialize_device_settings(
        use_cuda=gpu, local_rank=-1, use_amp=None
    )
    name = os.path.basename(model_name_or_path)

    # a) either from local dir
    if os.path.exists(model_name_or_path):
        model = BaseAdaptiveModel.load(
            load_dir=model_name_or_path, device=device, strict=strict
        )
        if task_type == "embeddings":
            processor = InferenceProcessor.load_from_dir(model_name_or_path)
        else:
            processor = Processor.load_from_dir(model_name_or_path)

    # b) or from remote transformers model hub
    else:
        logger.info(
            f"Could not find `{model_name_or_path}` locally. Try to download from model hub ..."
        )
        if not task_type:
            raise ValueError(
                "Please specify the 'task_type' of the model you want to load from transformers. "
                "Valid options for arg `task_type`:"
                "'question_answering', 'embeddings', 'text_classification', 'ner'"
            )

        model = AdaptiveModel.convert_from_transformers(
            model_name_or_path, device, task_type
        )
        config = AutoConfig.from_pretrained(model_name_or_path)
        tokenizer = Tokenizer.load(model_name_or_path)

        # TODO infer task_type automatically from config (if possible)
        if task_type == "question_answering":
            processor = SquadProcessor(
                tokenizer=tokenizer,
                max_seq_len=max_seq_len,
                label_list=["start_token", "end_token"],
                metric="squad",
                data_dir="data",
                doc_stride=doc_stride,
            )
        elif task_type == "embeddings":
            processor = InferenceProcessor(tokenizer=tokenizer, max_seq_len=max_seq_len)

        elif task_type == "text_classification":
            label_list = list(config.id2label[id] for id in range(len(config.id2label)))
            processor = TextClassificationProcessor(
                tokenizer=tokenizer,
                max_seq_len=max_seq_len,
                data_dir="data",
                label_list=label_list,
                label_column_name="label",
                metric="acc",
                quote_char='"',
            )
        elif task_type == "ner":
            label_list = list(config.label2id.keys())
            processor = NERProcessor(
                tokenizer=tokenizer,
                max_seq_len=max_seq_len,
                data_dir="data",
                metric="seq_f1",
                label_list=label_list,
            )
        else:
            raise ValueError(
                f"`task_type` {task_type} is not supported yet. "
                f"Valid options for arg `task_type`: 'question_answering', "
                f"'embeddings', 'text_classification', 'ner'"
            )

    return cls(
        model,
        processor,
        task_type=task_type,
        batch_size=batch_size,
        gpu=gpu,
        name=name,
        return_class_probs=return_class_probs,
        extraction_strategy=extraction_strategy,
        extraction_layer=extraction_layer,
    )