import json


def load_nli_jsonl(inp_file: str, language: str) -> list:
    """
    Load entailment sentence pairs from XNLI/MNLI jsonl file.

    Examples:
    ::
        # XNLI example
        {"annotator_labels":
            ["contradiction", "contradiction", "contradiction", "contradiction", "contradiction"],
         "genre": "facetoface",
         "gold_label": "contradiction",
         "language": "en",
         "match": "True",
         "pairID": "2",
         "promptID": "1",
         "sentence1": "And he said, Mama, I'm home.",
         "sentence1_tokenized": "And he said , Mama , I 'm home .",
         "sentence2": "He didn't say a word.",
         "sentence2_tokenized": "He didn 't say a word ."}

         # MNLI example
         {"annotator_labels":
             ["neutral", "entailment", "neutral", "neutral", "neutral"],
          "genre": "slate",
          "gold_label": "neutral",
          "pairID": "63735n",
          "promptID": "63735",
          "sentence1": "The new rights are nice enough",
          "sentence1_binary_parse": "( ( The ( new rights ) ) ( are ( nice enough ) ) )",
          "sentence1_parse": "(ROOT (S (NP (DT The) (JJ new) (NNS rights)) (VP (VBP are)
                             (ADJP (JJ nice) (RB enough)))))",
          "sentence2": "Everyone really likes the newest benefits ",
          "sentence2_binary_parse": "( Everyone ( really ( likes ( the ( newest benefits ) ) ) ) )",
          "sentence2_parse": "(ROOT (S (NP (NN Everyone)) (VP (ADVP (RB really)) (VBZ likes)
                             (NP (DT the) (JJS newest) (NNS benefits)))))"}
    """
    res = []
    with open(inp_file, "r") as f:
        for line in f:
            example = json.loads(line)
            # XNLI contains 15 different languages and we only use English examples
            if "language" in example and example["language"] != language:
                continue
            # use entailment paris to do training
            if example["gold_label"] != "entailment":
                continue
            res.append(example)

    return res