<gh_stars>1-10
import json
import logging
import os
from functools import partial
from pathlib import Path
from typing import Callable, Iterable, List

import click
import numpy as np
import pandas as pd
from rouge import rouge_score
from simpletransformers.seq2seq import Seq2SeqArgs, Seq2SeqModel
from sklearn.model_selection import train_test_split
from transformers import BartTokenizerFast

logging.basicConfig(level=logging.INFO)

# os.environ["TOKENIZERS_PARALLELISM"] = "false"


def precision(labels: Iterable, preds: Iterable):
    unique_labels = set(labels)
    unique_preds = set(preds)
    precision = len(unique_preds.intersection(unique_labels)) / len(
        unique_preds
    )
    return precision


def recall(labels: Iterable, preds: Iterable):
    unique_labels = set(labels)
    unique_preds = set(preds)
    recall = len(unique_labels.intersection(unique_preds)) / len(unique_labels)
    return recall


def flatten(
    in_list: Iterable[Iterable],
) -> List:
    return [item for sublist in in_list for item in sublist]


def test_print(labels, preds):
    for label, pred in zip(labels, preds):
        print(label, "|", pred)
        print("\n")


def macro_precision(labels: List[str], preds: List[List[str]]):
    preds_macro = flatten(preds)
    return precision(labels, preds_macro)


def macro_recall(labels: List[str], preds: List[List[str]]):
    preds_macro = flatten(preds)
    return recall(labels, preds_macro)


def micro_precision(labels: List[str], preds: List[List[str]]):
    precisions = []
    for labels_single, preds_single in zip(labels, preds):
        if type(labels_single) != str:
            raise ValueError
        precisions.append(precision([labels_single], preds_single))
    return np.mean(precisions)


def micro_recall(labels: List[str], preds: List[List[str]]):
    recalls = []
    for labels_single, preds_single in zip(labels, preds):
        recalls.append(recall([labels_single], preds_single))
    return np.mean(recalls)


def max_rouge(labels: List[str], preds: List[List[str]], n=2):
    out_f1s = []
    out_precisions = []
    out_recalls = []
    for labels_single, preds_single in zip(labels, preds):
        single_scores = []
        for pred in preds_single:
            single_scores.append(
                rouge_score.rouge_n([labels_single], [pred], n=n)
            )
        max_f1 = max([x["f"] for x in single_scores])
        max_precision = max([x["p"] for x in single_scores])
        max_recall = max([x["r"] for x in single_scores])
        out_f1s.append(max_f1)
        out_precisions.append(max_precision)
        out_recalls.append(max_recall)

    return {
        "f1": np.mean(out_f1s),
        "precision": np.mean(out_precisions),
        "recall": np.mean(out_recalls),
    }


def _load_ls_json(input_file):
    with open(input_file, "r") as f:
        data = json.load(f)
    pairs = []
    for el in data:
        premise = el["data"]["premise"]
        for annot in el["annotations"]:
            for result in annot["result"]:
                hypotheses = result["value"]["text"]
                for hypothesis in hypotheses:
                    pairs.append([premise, hypothesis])
    data_train = pd.DataFrame(pairs, columns=["input_text", "target_text"])
    return data_train


@click.command()
@click.option("--input-file", type=click.Path(exists=True), required=True)
@click.option(
    "--input-format",
    type=click.Choice(["ls-json", "st-csv"]),
    default="ls-json",
)
@click.option(
    "--positives-only",
    type=bool,
    help="Only use examples with at least one output tuple",
)
@click.option(
    "--output-dir",
    help="folder to store the trained model and training artifacts",
    type=click.Path(),
    required=True,
)
@click.option("--model-name", type=str, default="bart")
@click.option("--model-path", type=str, default="facebook/bart-large")
@click.option("--do-sample", type=bool, default=True)
@click.option("--num-beams", type=int, default=None)
@click.option("--top-k", type=int, default=10)
@click.option(
    "--max-length", type=int, default=64, help="max output sequence length"
)
@click.option(
    "--max-seq-length", type=int, default=512, help="max input sequence length"
)
@click.option("--length-penalty", type=float, default=2.0)
@click.option("--val-size", type=float, default=0.2)
@click.option("--num-train-epochs", type=int, default=5)
@click.option("--num-return-sequences", type=int, default=2)
@click.option("--subset-data-num-return-sequences", type=bool, default=False)
@click.option("--subset-data-max-seq-length", type=bool, default=True)
@click.option("--repetition-penalty", type=float, default=1000)
@click.option("--early-stopping-metric", type=str, default="eval_loss")
@click.option(
    "--early-stopping-delta",
    type=float,
    help=(
        "The improvement over best_eval_loss necessary "
        "to count as a better checkpoint."
    ),
    default=0.01,
)
@click.option(
    "--early-stopping-patience",
    type=int,
    default=3,
    help=(
        "Terminate training after this many evaluations without an "
        "improvement in the evaluation metric greater than "
        "early_stopping_delta."
    ),
)
@click.option(
    "--evaluate-during-training",
    type=bool,
    default=True,
)
@click.option(
    "--use-cuda",
    type=bool,
    default=True,
)
def train_seq2seq(
    input_file,
    input_format: str,
    positives_only: bool,
    output_dir,
    model_name: str,
    model_path: str,
    do_sample: bool,
    num_beams: int,
    top_k: int,
    max_length: int,
    max_seq_length: int,
    length_penalty: float,
    val_size: float,
    num_train_epochs: int,
    num_return_sequences: int,
    subset_data_num_return_sequences: bool,
    subset_data_max_seq_length: bool,
    repetition_penalty: float,
    early_stopping_metric: str,
    early_stopping_delta: float,
    early_stopping_patience: int,
    evaluate_during_training: bool,
    use_cuda: bool,
):
    max_rouge("test sequence here", ["test sequence here"], n=2)
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True)

    if input_format == "ls-json":
        data_train = _load_ls_json(input_file)
    elif input_format == "st-csv":
        data_train = pd.read_csv(input_file)

    data_train = data_train[data_train.target_text != "ambiguous"]
    data_train = data_train[~data_train.target_text.isna()]
    logging.info(f"Loaded {len(data_train)} pairs")

    if positives_only:
        logging.info("Filtering to only positive examples")
        data_train = data_train[data_train["target_text"] != "none"]
        logging.info(f"Filtered to {len(data_train)} pairs")

    if subset_data_num_return_sequences:
        logging.info(
            f"Filtering to only data with num_return_sequences == {num_return_sequences}"
        )
        output_counts = data_train.groupby("input_text").agg(
            {"target_text": "count"}
        )["target_text"]
        accept_inputs = output_counts[
            output_counts == num_return_sequences
        ].index.values
        data_train = data_train[data_train["input_text"].isin(accept_inputs)]
        logging.info(f"Filtered to {len(data_train)} pairs")

    model_args = Seq2SeqArgs()

    model_args.do_sample = do_sample
    model_args.evaluate_during_training = evaluate_during_training
    model_args.evaluate_generated_text = True
    model_args.evaluate_during_training_verbose = True
    model_args.num_beams = num_beams
    model_args.num_return_sequences = num_return_sequences
    model_args.repetition_penalty = float(repetition_penalty)
    model_args.early_stopping_metric = early_stopping_metric
    model_args.early_stopping_delta = early_stopping_delta
    model_args.early_stopping_patience = early_stopping_patience
    model_args.top_k = top_k
    model_args.max_length = max_length
    model_args.max_seq_length = max_seq_length
    model_args.length_penalty = length_penalty

    model_args.num_train_epochs = num_train_epochs
    model_args.train_batch_size = 2
    model_args.gradient_accumulation_steps = 4
    model_args.overwrite_output_dir = True
    model_args.save_model_every_epoch = False
    model_args.use_multiprocessing = (
        False  # doesn't work since we updated the slurm config
    )
    model_args.use_multiprocessing_for_evaluation = (
        False  # doesn't work since we updated the slurm config
    )
    model_args.reprocess_input_data = True
    model_args.no_cache = True
    model_args.no_save = True

    model = Seq2SeqModel(
        encoder_decoder_type=model_name,
        encoder_decoder_name=model_path,
        args=model_args,
        use_cuda=use_cuda,
    )

    if subset_data_max_seq_length:
        tokenizer = BartTokenizerFast.from_pretrained(model_path)
        tokenized = tokenizer.batch_encode_plus(data_train.input_text.tolist())
        input_lens = pd.Series([len(x) for x in tokenized["input_ids"]])
        data_train = data_train[input_lens <= max_seq_length]
        logging.info(f"Trimmed to {len(data_train)} pairs")

    # TODO this should have the option to split on the input_text
    data_train, data_val = train_test_split(data_train, test_size=val_size)
    logging.info(f"Split to {len(data_train)} train and {len(data_val)} val")

    # data_train["target_text"] = data_train["target_text"] + " [EOS]"
    train_file = output_dir / "train.csv"
    logging.info(f"Saving training data to {train_file}")
    data_train.to_csv(train_file, index=False)

    if num_return_sequences > 1:
        metric_kwargs = {
            "test_print": test_print,
            "macro_precision": macro_precision,
            "micro_precision": micro_precision,
            "macro_recall": macro_recall,
            "micro_recall": micro_recall,
            "rouge-1": lambda labels, preds: max_rouge(labels, preds, n=1),
            "rouge-2": lambda labels, preds: max_rouge(labels, preds, n=2),
            "rouge-3": lambda labels, preds: max_rouge(labels, preds, n=3),
        }
    elif num_return_sequences == 1:
        metric_kwargs = {
            "test_print": test_print,
            "exact_match": lambda labels, preds: np.mean(
                [l == p for l, p in zip(labels, preds)]
            ),
            "rouge-1": lambda labels, preds: rouge_score.rouge_n(
                labels, preds, n=1
            ),
            "rouge-2": lambda labels, preds: rouge_score.rouge_n(
                labels, preds, n=2
            ),
            "rouge-3": lambda labels, preds: rouge_score.rouge_n(
                labels, preds, n=3
            ),
        }

    model.train_model(
        data_train,
        eval_data=data_train.sample(50),
        **metric_kwargs,
    )
    logging.info("Evaluating model")
    eval_results = model.eval_model(data_train, **metric_kwargs)
    logging.info(eval_results)
    eval_file = output_dir / "eval.json"
    logging.info(f"Saving evaluation results to {eval_file}")
    with open(eval_file, "w") as f:
        json.dump(eval_results, f)
    model_file = output_dir / "model.ckpt"
    logging.info("Saving model to %s", model_file)
    model.save_model(model_file)


if __name__ == "__main__":
    train_seq2seq()

import json
import logging
import os
from functools import partial
from pathlib import Path
from typing import Callable, Iterable, List

import click
import numpy as np
import pandas as pd
from rouge import rouge_score
from simpletransformers.seq2seq import Seq2SeqArgs, Seq2SeqModel
from sklearn.model_selection import train_test_split
from transformers import BartTokenizerFast

logging.basicConfig(level=logging.INFO)

# os.environ["TOKENIZERS_PARALLELISM"] = "false"


def precision(labels: Iterable, preds: Iterable):
    unique_labels = set(labels)
    unique_preds = set(preds)
    precision = len(unique_preds.intersection(unique_labels)) / len(
        unique_preds
    )
    return precision


def recall(labels: Iterable, preds: Iterable):
    unique_labels = set(labels)
    unique_preds = set(preds)
    recall = len(unique_labels.intersection(unique_preds)) / len(unique_labels)
    return recall


def flatten(
    in_list: Iterable[Iterable],
) -> List:
    return [item for sublist in in_list for item in sublist]


def test_print(labels, preds):
    for label, pred in zip(labels, preds):
        print(label, "|", pred)
        print("\n")


def macro_precision(labels: List[str], preds: List[List[str]]):
    preds_macro = flatten(preds)
    return precision(labels, preds_macro)


def macro_recall(labels: List[str], preds: List[List[str]]):
    preds_macro = flatten(preds)
    return recall(labels, preds_macro)


def micro_precision(labels: List[str], preds: List[List[str]]):
    precisions = []
    for labels_single, preds_single in zip(labels, preds):
        if type(labels_single) != str:
            raise ValueError
        precisions.append(precision([labels_single], preds_single))
    return np.mean(precisions)


def micro_recall(labels: List[str], preds: List[List[str]]):
    recalls = []
    for labels_single, preds_single in zip(labels, preds):
        recalls.append(recall([labels_single], preds_single))
    return np.mean(recalls)


def max_rouge(labels: List[str], preds: List[List[str]], n=2):
    out_f1s = []
    out_precisions = []
    out_recalls = []
    for labels_single, preds_single in zip(labels, preds):
        single_scores = []
        for pred in preds_single:
            single_scores.append(
                rouge_score.rouge_n([labels_single], [pred], n=n)
            )
        max_f1 = max([x["f"] for x in single_scores])
        max_precision = max([x["p"] for x in single_scores])
        max_recall = max([x["r"] for x in single_scores])
        out_f1s.append(max_f1)
        out_precisions.append(max_precision)
        out_recalls.append(max_recall)

    return {
        "f1": np.mean(out_f1s),
        "precision": np.mean(out_precisions),
        "recall": np.mean(out_recalls),
    }


def _load_ls_json(input_file):
    with open(input_file, "r") as f:
        data = json.load(f)
    pairs = []
    for el in data:
        premise = el["data"]["premise"]
        for annot in el["annotations"]:
            for result in annot["result"]:
                hypotheses = result["value"]["text"]
                for hypothesis in hypotheses:
                    pairs.append([premise, hypothesis])
    data_train = pd.DataFrame(pairs, columns=["input_text", "target_text"])
    return data_train


@click.command()
@click.option("--input-file", type=click.Path(exists=True), required=True)
@click.option(
    "--input-format",
    type=click.Choice(["ls-json", "st-csv"]),
    default="ls-json",
)
@click.option(
    "--positives-only",
    type=bool,
    help="Only use examples with at least one output tuple",
)
@click.option(
    "--output-dir",
    help="folder to store the trained model and training artifacts",
    type=click.Path(),
    required=True,
)
@click.option("--model-name", type=str, default="bart")
@click.option("--model-path", type=str, default="facebook/bart-large")
@click.option("--do-sample", type=bool, default=True)
@click.option("--num-beams", type=int, default=None)
@click.option("--top-k", type=int, default=10)
@click.option(
    "--max-length", type=int, default=64, help="max output sequence length"
)
@click.option(
    "--max-seq-length", type=int, default=512, help="max input sequence length"
)
@click.option("--length-penalty", type=float, default=2.0)
@click.option("--val-size", type=float, default=0.2)
@click.option("--num-train-epochs", type=int, default=5)
@click.option("--num-return-sequences", type=int, default=2)
@click.option("--subset-data-num-return-sequences", type=bool, default=False)
@click.option("--subset-data-max-seq-length", type=bool, default=True)
@click.option("--repetition-penalty", type=float, default=1000)
@click.option("--early-stopping-metric", type=str, default="eval_loss")
@click.option(
    "--early-stopping-delta",
    type=float,
    help=(
        "The improvement over best_eval_loss necessary "
        "to count as a better checkpoint."
    ),
    default=0.01,
)
@click.option(
    "--early-stopping-patience",
    type=int,
    default=3,
    help=(
        "Terminate training after this many evaluations without an "
        "improvement in the evaluation metric greater than "
        "early_stopping_delta."
    ),
)
@click.option(
    "--evaluate-during-training",
    type=bool,
    default=True,
)
@click.option(
    "--use-cuda",
    type=bool,
    default=True,
)
def train_seq2seq(
    input_file,
    input_format: str,
    positives_only: bool,
    output_dir,
    model_name: str,
    model_path: str,
    do_sample: bool,
    num_beams: int,
    top_k: int,
    max_length: int,
    max_seq_length: int,
    length_penalty: float,
    val_size: float,
    num_train_epochs: int,
    num_return_sequences: int,
    subset_data_num_return_sequences: bool,
    subset_data_max_seq_length: bool,
    repetition_penalty: float,
    early_stopping_metric: str,
    early_stopping_delta: float,
    early_stopping_patience: int,
    evaluate_during_training: bool,
    use_cuda: bool,
):
    max_rouge("test sequence here", ["test sequence here"], n=2)
    output_dir = Path(output_dir)
    output_dir.mkdir(exist_ok=True)

    if input_format == "ls-json":
        data_train = _load_ls_json(input_file)
    elif input_format == "st-csv":
        data_train = pd.read_csv(input_file)

    data_train = data_train[data_train.target_text != "ambiguous"]
    data_train = data_train[~data_train.target_text.isna()]
    logging.info(f"Loaded {len(data_train)} pairs")

    if positives_only:
        logging.info("Filtering to only positive examples")
        data_train = data_train[data_train["target_text"] != "none"]
        logging.info(f"Filtered to {len(data_train)} pairs")

    if subset_data_num_return_sequences:
        logging.info(
            f"Filtering to only data with num_return_sequences == {num_return_sequences}"
        )
        output_counts = data_train.groupby("input_text").agg(
            {"target_text": "count"}
        )["target_text"]
        accept_inputs = output_counts[
            output_counts == num_return_sequences
        ].index.values
        data_train = data_train[data_train["input_text"].isin(accept_inputs)]
        logging.info(f"Filtered to {len(data_train)} pairs")

    model_args = Seq2SeqArgs()

    model_args.do_sample = do_sample
    model_args.evaluate_during_training = evaluate_during_training
    model_args.evaluate_generated_text = True
    model_args.evaluate_during_training_verbose = True
    model_args.num_beams = num_beams
    model_args.num_return_sequences = num_return_sequences
    model_args.repetition_penalty = float(repetition_penalty)
    model_args.early_stopping_metric = early_stopping_metric
    model_args.early_stopping_delta = early_stopping_delta
    model_args.early_stopping_patience = early_stopping_patience
    model_args.top_k = top_k
    model_args.max_length = max_length
    model_args.max_seq_length = max_seq_length
    model_args.length_penalty = length_penalty

    model_args.num_train_epochs = num_train_epochs
    model_args.train_batch_size = 2
    model_args.gradient_accumulation_steps = 4
    model_args.overwrite_output_dir = True
    model_args.save_model_every_epoch = False
    model_args.use_multiprocessing = (
        False  # doesn't work since we updated the slurm config
    )
    model_args.use_multiprocessing_for_evaluation = (
        False  # doesn't work since we updated the slurm config
    )
    model_args.reprocess_input_data = True
    model_args.no_cache = True
    model_args.no_save = True

    model = Seq2SeqModel(
        encoder_decoder_type=model_name,
        encoder_decoder_name=model_path,
        args=model_args,
        use_cuda=use_cuda,
    )

    if subset_data_max_seq_length:
        tokenizer = BartTokenizerFast.from_pretrained(model_path)
        tokenized = tokenizer.batch_encode_plus(data_train.input_text.tolist())
        input_lens = pd.Series([len(x) for x in tokenized["input_ids"]])
        data_train = data_train[input_lens <= max_seq_length]
        logging.info(f"Trimmed to {len(data_train)} pairs")

    # TODO this should have the option to split on the input_text
    data_train, data_val = train_test_split(data_train, test_size=val_size)
    logging.info(f"Split to {len(data_train)} train and {len(data_val)} val")

    # data_train["target_text"] = data_train["target_text"] + " [EOS]"
    train_file = output_dir / "train.csv"
    logging.info(f"Saving training data to {train_file}")
    data_train.to_csv(train_file, index=False)

    if num_return_sequences > 1:
        metric_kwargs = {
            "test_print": test_print,
            "macro_precision": macro_precision,
            "micro_precision": micro_precision,
            "macro_recall": macro_recall,
            "micro_recall": micro_recall,
            "rouge-1": lambda labels, preds: max_rouge(labels, preds, n=1),
            "rouge-2": lambda labels, preds: max_rouge(labels, preds, n=2),
            "rouge-3": lambda labels, preds: max_rouge(labels, preds, n=3),
        }
    elif num_return_sequences == 1:
        metric_kwargs = {
            "test_print": test_print,
            "exact_match": lambda labels, preds: np.mean(
                [l == p for l, p in zip(labels, preds)]
            ),
            "rouge-1": lambda labels, preds: rouge_score.rouge_n(
                labels, preds, n=1
            ),
            "rouge-2": lambda labels, preds: rouge_score.rouge_n(
                labels, preds, n=2
            ),
            "rouge-3": lambda labels, preds: rouge_score.rouge_n(
                labels, preds, n=3
            ),
        }

    model.train_model(
        data_train,
        eval_data=data_train.sample(50),
        **metric_kwargs,
    )
    logging.info("Evaluating model")
    eval_results = model.eval_model(data_train, **metric_kwargs)
    logging.info(eval_results)
    eval_file = output_dir / "eval.json"
    logging.info(f"Saving evaluation results to {eval_file}")
    with open(eval_file, "w") as f:
        json.dump(eval_results, f)
    model_file = output_dir / "model.ckpt"
    logging.info("Saving model to %s", model_file)
    model.save_model(model_file)


if __name__ == "__main__":
    train_seq2seq()
# os.environ["TOKENIZERS_PARALLELISM"] = "false" # doesn't work since we updated the slurm config # doesn't work since we updated the slurm config # TODO this should have the option to split on the input_text # data_train["target_text"] = data_train["target_text"] + " [EOS]"