examples/research_projects/jax-projects/big_bird/evaluate.py/main
def main():
    dataset = load_from_disk("natural-questions-validation")
    dataset = dataset.map(format_dataset).remove_columns(["annotations", "document", "id"])
    print(dataset)

    short_validation_dataset = dataset.filter(lambda x: (len(x["question"]) + len(x["context"])) < 4 * 4096)
    short_validation_dataset = short_validation_dataset.filter(lambda x: x["category"] != "null")
    short_validation_dataset

    model_id = "vasudevgupta/flax-bigbird-natural-questions"
    model = FlaxBigBirdForNaturalQuestions.from_pretrained(model_id)
    tokenizer = BigBirdTokenizerFast.from_pretrained(model_id)

    @jax.jit
    def forward(*args, **kwargs):
        start_logits, end_logits, pooled_logits = model(*args, **kwargs)
        return start_logits, end_logits, jnp.argmax(pooled_logits, axis=-1)

    def evaluate(example):
        # encode question and context so that they are seperated by a tokenizer.sep_token and cut at max_length
        inputs = tokenizer(
            example["question"],
            example["context"],
            return_tensors="np",
            max_length=4096,
            padding="max_length",
            truncation=True,
        )

        start_scores, end_scores, category = forward(**inputs)

        predicted_category = CATEGORY_MAPPING[category.item()]

        example["targets"] = example["long"] + example["short"]
        if example["category"] in ["yes", "no", "null"]:
            example["targets"] = [example["category"]]
        example["has_tgt"] = example["category"] != "null"
        # Now target can be: "yes", "no", "null", "list of long & short answers"

        if predicted_category in ["yes", "no", "null"]:
            example["output"] = [predicted_category]
            example["match"] = example["output"] == example["targets"]
            example["has_pred"] = predicted_category != "null"
            return example

        max_size = 38 if predicted_category == "short" else 1024
        start_score, end_score = get_best_valid_start_end_idx(
            start_scores[0], end_scores[0], top_k=8, max_size=max_size
        )

        input_ids = inputs["input_ids"][0].tolist()
        example["output"] = [tokenizer.decode(input_ids[start_score : end_score + 1])]

        answers = expand_to_aliases(example["targets"], make_sub_answers=True)
        predictions = expand_to_aliases(example["output"])

        # some preprocessing to both prediction and answer
        answers = set(["".join(a.split()) for a in answers])
        predictions = set(["".join(p.split()) for p in predictions])
        predictions = set([s for s in predictions if s not in ["``", "''", "`", "'"]])

        # if there is a common element, it's a exact match
        example["match"] = len(list(answers & predictions)) > 0
        example["has_pred"] = predicted_category != "null" and len(predictions) > 0

        return example

    short_validation_dataset = short_validation_dataset.map(evaluate)

    total = len(short_validation_dataset)
    matched = len(short_validation_dataset.filter(lambda x: x["match"] == 1))
    print("EM score:", (matched / total) * 100, "%")