RunNLPfinal.py

import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional
from transformers import Trainer
from transformers.trainer_utils import PredictionOutput

import datasets
from datasets import load_dataset, load_metric

import transformers
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    EvalPrediction,
    HfArgumentParser,
    PreTrainedTokenizerFast,
    TrainingArguments,
    default_data_collator,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint
from transformers.utils import check_min_version
from transformers.utils.versions import require_version
from utilsNLPfinal import postprocesspredictions,helper_function_process,helper_function

check_min_version("4.9.0")

require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/question-answering/requirements.txt")

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:


    model_name_or_path: str = field(
        default=None,
        metadata={"help": "Either  Path to pretrained model or model identifier at huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path"}
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Path(of directory) to store the pretrained models from huggingface.co"},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use."},
    )
    use_auth_token: bool = field(
        default=False,
        metadata={ "help": "Defaults to token generated by `transformers-cli login` (for private models)"},
    )


@dataclass
class DataTrainingArguments:


    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (from the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use."}
    )
    train_file: Optional[str] = field(default=None, metadata={"help": "Input training data file(.txt)."})
    validation_file: Optional[str] = field(
        default=None,
        metadata={"help": "(Optional) input evaluation data file to evaluate perplexity (.txt)."},
    )
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "(Optional) input test data file to evaluate perplexity (.txt)."},
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={"help": "No. of processes for preprocessing."},
    )
    do_not_use_token_type_ids: bool = field(
        default=False, metadata={"help": "If true, don't use token_type_ids."}
    )
    max_seq_length: int = field(
        default=384,
        metadata={
            "help": "Max total input sequence length after tokenization. Longer ones will be truncated and shorter ones will be padded"
        },
    )
    pad_to_max_length: bool = field(
        default=True,
        metadata={
            "help": "Should all samples be padded to `max_seq_length`. "
                    "If False, will pad the samples dynamically when batching to the maximum length in the batch (which can "
                    "be faster on GPU but will be slower on TPU)."
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging, truncate the number of training examples"
                    "value if set."
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging, truncate the number of evaluation examples"
                    "value if set."
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": "For debugging, truncate the number of prediction examples"
                    "value if set."
        },
    )
    version_2_with_negative: bool = field(
        default=False, metadata={"help": "If true, some examples don't have answers."}
    )
    null_score_diff_threshold: float = field(
        default=0.0,
        metadata={
            "help": "The threshold for null answer: if the best answer has a score that is less than "
                    "the score of null answer minus this threshold, the null answer is selected. "
                    "Only used when `version_2_with_negative=True`."
        },
    )
    doc_stride: int = field(
        default=128,
        metadata={"help": "For splitting up a long document into chunks, set stride to take between chunks."},
    )
    n_best_size: int = field(
        default=20,
        metadata={"help": "The total number of n-best predictions to generate."},
    )
    max_answer_length: int = field(
        default=30,
        metadata={
            "help": "The maximum length of an answer that can be generated. As start and end predictions are not conditioned on one another."
        },
    )

    def __post_init__(self):
        if self.dataset_name is not None or self.train_file is not None or self.validation_file is not None or self.test_file is not None:
            if self.train_file is None:
                pass
            else:
                extension = self.train_file.split(".")[-1]
                assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
            if self.validation_file is None:
                pass
            else:
                extension = self.validation_file.split(".")[-1]
                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
            if self.test_file is None:
                return
            extension = self.test_file.split(".")[-1]
            assert extension in ["csv", "json"], "`test_file` should be a csv or a json file."
        else:
            raise ValueError("Need either a dataset name or a training/validation file/test_file.")


class QandATrainer(Trainer):
    def __init__(self, *args, eval_examples=None, post_process_function=None, **kwargs):
        super().__init__(*args, **kwargs)
        self.eval_examples = eval_examples
        self.post_process_function = post_process_function

    def predict(self, predict_dataset, predict_examples, ignore_keys=None, metric_key_prefix: str = "test"):
        compute_metrics, eval_loop, predict_dataloader = self.predict_initialise(predict_dataset)
        try:
            output = eval_loop(
                predict_dataloader,
                description="Prediction",
                prediction_loss_only=True if compute_metrics is None else None,
                ignore_keys=ignore_keys,
            )
        finally:
            self.compute_metrics = compute_metrics

        if self.post_process_function is None:
            return output
        if self.compute_metrics is None:
            return output
        predictions = self.post_process_function(predict_examples, predict_dataset, output.predictions, "predict")
        metrics = self.compute_metrics(predictions)
        for key in list(metrics.keys()):
            if key.startswith(f"{metric_key_prefix}_"):
                continue
            metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
        self.log(metrics)

        return PredictionOutput(predictions=predictions.predictions, label_ids=predictions.label_ids, metrics=metrics)

    def predict_initialise(self, predict_dataset):
        data_load = self.get_test_dataloader(predict_dataset)
        compute_metrics = self.compute_metrics
        self.compute_metrics = None
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        return compute_metrics, eval_loop, data_load

    def evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metric_key_prefix: str = "eval"):
        compute_metrics, eval_dataloader, eval_dataset, eval_examples, eval_loop = self.evaulate_init_modelmodel(
            eval_dataset, eval_examples)
        try:
            output = eval_loop(
                eval_dataloader,
                description="Evaluation",
                prediction_loss_only=True if compute_metrics is None else None,
                ignore_keys=ignore_keys,
            )
        finally:
            self.compute_metrics = compute_metrics
        if self.post_process_function is None or self.compute_metrics is None:
            metrics = {}
        else:
            eval_preds = self.post_process_function(eval_examples, eval_dataset, output.predictions)
            metrics = self.compute_metrics(eval_preds)
            for key in list(metrics.keys()):
                if not key.startswith(f"{metric_key_prefix}_"):
                    metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

            self.log(metrics)

        self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, metrics)
        return metrics

    def evaulate_init_modelmodel(self, eval_dataset, eval_examples):
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        eval_examples = eval_examples if eval_examples is not None else self.eval_examples
        compute_metrics = self.compute_metrics
        self.compute_metrics = None
        eval_loop = self.evaluation_loop if not self.args.use_legacy_prediction_loop else self.prediction_loop
        return compute_metrics, eval_dataloader, eval_dataset, eval_examples, eval_loop


def main():


    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
    if len(sys.argv) != 2 or not sys.argv[1].endswith(".json"):
        model_args, data_args, args_train = parser.parse_args_into_dataclasses()
    else:
        model_args, data_args, args_train = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        handlers=[logging.StreamHandler(sys.stdout)],
    )

    log_level = args_train.get_process_log_level()
    logger.setLevel(log_level)
    datasets.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.set_verbosity(log_level)
    transformers.utils.logging.enable_default_handler()
    transformers.utils.logging.enable_explicit_format()

    logger.warning(
        f"Process rank: {args_train.local_rank}, device: {args_train.device}, n_gpu: {args_train.n_gpu}"
        + f"distributed training: {bool(args_train.local_rank != -1)}, 16-bits training: {args_train.fp16}"
    )
    logger.info(f"Training/evaluation parameters {args_train}")

    last_checkpoint = None
    if not os.path.isdir(args_train.output_dir) or not args_train.do_train or args_train.overwrite_output_dir:
        pass
    else:
        last_checkpoint = get_last_checkpoint(args_train.output_dir)
        if last_checkpoint is None:
            if len(os.listdir(args_train.output_dir)) > 0:
                raise ValueError(
                    f"Output directory ({args_train.output_dir}) already exists and is not empty. "
                    "Use --overwrite_output_dir to overcome."
                )
            elif last_checkpoint is not None and args_train.resume_from_checkpoint is None:
                logger.info(
                    f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                    "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
                )
        elif last_checkpoint is not None and args_train.resume_from_checkpoint is None:
            logger.info(
                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
            )

    set_seed(args_train.seed)

    if data_args.dataset_name is None:
        data_files = {}
        if data_args.train_file is None:
            pass
        else:
            data_files["train"] = data_args.train_file
            extension = data_args.train_file.split(".")[-1]

        if data_args.validation_file is None:
            pass
        else:
            data_files["validation"] = data_args.validation_file
            extension = data_args.validation_file.split(".")[-1]
        if data_args.test_file is None:
            pass
        else:
            data_files["test"] = data_args.test_file
            extension = data_args.test_file.split(".")[-1]
        raw_datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)  # field="data",
    else:
        raw_datasets = load_dataset(
            data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir
        )

    config = AutoConfig.from_pretrained(
        model_args.model_name_or_path if not model_args.config_name else model_args.config_name,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=None if not model_args.use_auth_token else True,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path if not model_args.tokenizer_name else model_args.tokenizer_name,
        cache_dir=model_args.cache_dir,
        use_fast=True,
        revision=model_args.model_revision,
        use_auth_token=None if not model_args.use_auth_token else True,
    )

    if not model_args.model_name_or_path:

        print("configure happening:", config)
        model = AutoModelForQuestionAnswering.from_config(config)
    else:
        model = AutoModelForQuestionAnswering.from_pretrained(
            model_args.model_name_or_path,
            from_tf=bool(".ckpt" in model_args.model_name_or_path),
            config=config,
            cache_dir=model_args.cache_dir,
            revision=model_args.model_revision,
            use_auth_token=None if not model_args.use_auth_token else True,
        )
    if isinstance(tokenizer, PreTrainedTokenizerFast):
        if args_train.do_train:
            column_names = raw_datasets["train"].column_names
        elif args_train.do_eval:
            column_names = raw_datasets["validation"].column_names
        else:
            column_names = raw_datasets["test"].column_names
        question_column_name = "question" if "question" in column_names else column_names[0]
        context_column_name = "context" if "context" in column_names else column_names[1]
        answer_column_name = "answers" if "answers" in column_names else column_names[2]

        pad_on_right = tokenizer.padding_side == "right"

        if data_args.max_seq_length <= tokenizer.model_max_length:
            pass
        else:
            logger.warning(
                f"The max_seq_length passed ({data_args.max_seq_length}) is larger than the maximum length for the"
                f"model ({tokenizer.model_max_length}). Using max_seq_length={tokenizer.model_max_length}."
            )
        max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)

        def prepare_train_features(examples):

            tokenized_examples = tokenizer(
                examples[context_column_name if not pad_on_right else question_column_name],
                examples[question_column_name if not pad_on_right else context_column_name],
                truncation="only_first" if not pad_on_right else "only_second",
                max_length=max_seq_length,
                stride=data_args.doc_stride,
                return_overflowing_tokens=True,
                return_offsets_mapping=True,
                padding=False if not data_args.pad_to_max_length else "max_length",
            )
            if not data_args.do_not_use_token_type_ids:
                pass
            else:
                if "token_type_ids" not in tokenized_examples:
                    return
                print("drop the token_type_ids.okay?")
                tokenized_examples.pop("token_type_ids")
                assert "token_type_ids" not in tokenized_examples
            sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
            offsetmapper = tokenized_examples.pop("offset_mapping")
            tokenized_examples["start_positions"] = []
            tokenized_examples["end_positions"] = []

            for num, offsets in enumerate(offsetmapper):
                input_ids = tokenized_examples["input_ids"][num]
                cls_index = input_ids.index(tokenizer.cls_token_id)
                sequence_ids = tokenized_examples.sequence_ids(num)
                sample_index = sample_mapping[num]
                answers = examples[answer_column_name][sample_index]
                if len(answers["answer_start"]) != 0:
                    start_char = answers["answer_start"][0]
                    end_char = start_char + len(answers["text"][0])

                    token_start_index = 0
                    while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
                        token_start_index += 1

                    token_end_index = len(input_ids) - 1
                    while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
                        token_end_index -= 1

                    if offsets[token_start_index][0] > start_char or offsets[token_end_index][1] < end_char:
                        tokenized_examples["start_positions"].append(cls_index)
                        tokenized_examples["end_positions"].append(cls_index)
                    else:
                        while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
                            token_start_index += 1
                        tokenized_examples["start_positions"].append(token_start_index - 1)
                        while offsets[token_end_index][1] >= end_char:
                            token_end_index -= 1
                        tokenized_examples["end_positions"].append(token_end_index + 1)
                else:
                    tokenized_examples["start_positions"].append(cls_index)
                    tokenized_examples["end_positions"].append(cls_index)

            return tokenized_examples

        if not args_train.do_train:
            pass
        else:
            if "train" in raw_datasets:
                pass
            else:
                raise ValueError("--do_train requires a train dataset")
            train_dataset = raw_datasets["train"]
            if data_args.max_train_samples is None:
                pass
            else:
                train_dataset = train_dataset.select(range(data_args.max_train_samples))
            with args_train.main_process_first(desc="train dataset map pre-processing"):
                train_dataset = train_dataset.map(
                    prepare_train_features,
                    batched=True,
                    num_proc=data_args.preprocessing_num_workers,
                    remove_columns=column_names,
                    load_from_cache_file=not data_args.overwrite_cache,
                    desc="Running tokenizer on train dataset",
                )
            if data_args.max_train_samples is None:
                return
            train_dataset = train_dataset.select(range(data_args.max_train_samples))

        # Validation preprocessing
        def prepare_validation_features(examples):

            tokenized_examples = tokenizer(
                examples[context_column_name if not pad_on_right else question_column_name],
                examples[question_column_name if not pad_on_right else context_column_name],
                truncation="only_first" if not pad_on_right else "only_second",
                max_length=max_seq_length,
                stride=data_args.doc_stride,
                return_overflowing_tokens=True,
                return_offsets_mapping=True,
                padding=False if not data_args.pad_to_max_length else "max_length",
            )
            if not data_args.do_not_use_token_type_ids:
                pass
            else:
                if "token_type_ids" not in tokenized_examples:
                    return
                print("drop token_type_ids!")
                tokenized_examples.pop("token_type_ids")
                assert "token_type_ids" not in tokenized_examples
            sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
            tokenized_examples["example_id"] = []

            helper_val_features(examples, sample_mapping, tokenized_examples)

            return tokenized_examples

        def helper_val_features(examples, sample_mapping, tokenized_examples):
            for i in range(len(tokenized_examples["input_ids"])):
                sequence_ids = tokenized_examples.sequence_ids(i)
                context_index = 0 if not pad_on_right else 1
                sample_index = sample_mapping[i]
                tokenized_examples["example_id"].append(examples["id"][sample_index])
                tokenized_examples["offset_mapping"][i] = []
                for k, o in enumerate(tokenized_examples["offset_mapping"][i]):
                    tokenized_examples["offset_mapping"][i].append((None if sequence_ids[k] != context_index else o))

        if not args_train.do_eval:
            pass
        else:
            if "validation" in raw_datasets:
                pass
            else:
                raise ValueError("--do_eval requires a validation dataset")
            eval_examples = raw_datasets["validation"]
            if data_args.max_eval_samples is None:
                pass
            else:
                eval_examples = eval_examples.select(range(data_args.max_eval_samples))
            with args_train.main_process_first(desc="validation dataset map pre-processing"):
                eval_dataset = eval_examples.map(
                    prepare_validation_features,
                    batched=True,
                    num_proc=data_args.preprocessing_num_workers,
                    remove_columns=column_names,
                    load_from_cache_file=not data_args.overwrite_cache,
                    desc="Running tokenizer on validation dataset",
                )
            if data_args.max_eval_samples is None:
                return
            eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))

        if not args_train.do_predict:
            pass
        else:
            if "test" in raw_datasets:
                pass
            else:
                raise ValueError("--do_predict requires a test dataset")
            predict_examples = raw_datasets["test"]
            if data_args.max_predict_samples is None:
                pass
            else:
                predict_examples = predict_examples.select(range(data_args.max_predict_samples))
            with args_train.main_process_first(desc="prediction dataset map pre-processing"):
                predict_dataset = predict_examples.map(
                    prepare_validation_features,
                    batched=True,
                    num_proc=data_args.preprocessing_num_workers,
                    remove_columns=column_names,
                    load_from_cache_file=not data_args.overwrite_cache,
                    desc="Running tokenizer on prediction dataset",
                )
            if data_args.max_predict_samples is None:
                return
            predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
        data_collator = (
            DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8 if args_train.fp16 else None)
            if not data_args.pad_to_max_length
            else default_data_collator
        )

        def post_processing_function(examples, features, predictions, stage="eval"):
            predictions = postprocesspredictions(
                examples=examples,
                features=features,
                predictions=predictions,
                version_2_with_negative=data_args.version_2_with_negative,
                n_best_size=data_args.n_best_size,
                max_answer_length=data_args.max_answer_length,
                null_score_diff_threshold=data_args.null_score_diff_threshold,
                output_dir=args_train.output_dir,
                log_level=log_level,
                prefix=stage,
            )
            if data_args.version_2_with_negative:
                formatted_predictions = [
                    {"id": k, "prediction_text": v, "no_answer_probability": 0.0} for k, v in predictions.items()
                ]
            else:
                formatted_predictions = [{"id": k, "prediction_text": v} for k, v in predictions.items()]

            references = [{"id": ex["id"], "answers": ex[answer_column_name]} for ex in examples]
            return EvalPrediction(predictions=formatted_predictions, label_ids=references)

        metric = load_metric("squad" if not data_args.version_2_with_negative else "squad_v2")

        def compute_metrics(p: EvalPrediction):
            return metric.compute(predictions=p.predictions, references=p.label_ids)

        trainer = QandATrainer(
            model=model,
            args=args_train,
            train_dataset=None if not args_train.do_train else train_dataset,
            eval_dataset=None if not args_train.do_eval else eval_dataset,
            eval_examples=None if not args_train.do_eval else eval_examples,
            tokenizer=tokenizer,
            data_collator=data_collator,
            post_process_function=post_processing_function,
            compute_metrics=compute_metrics,
        )

        if not args_train.do_train:
            pass
        else:
            checkpoint = None
            if args_train.resume_from_checkpoint is not None:
                checkpoint = args_train.resume_from_checkpoint
            elif last_checkpoint is not None:
                checkpoint = last_checkpoint
            train_result = trainer.train(resume_from_checkpoint=checkpoint)
            trainer.save_model()  # Saves the tokenizer too for easy upload

            metricstrain = train_result.metrics
            max_train_samples = (
                len(train_dataset) if data_args.max_train_samples is None else data_args.max_train_samples
            )
            metricstrain["train_samples"] = min(max_train_samples, len(train_dataset))

            trainer.log_metrics("train", metricstrain)
            trainer.save_metrics("train", metricstrain)
            trainer.save_state()

        if not args_train.do_eval:
            pass
        else:
            logger.info(" Evaluate ")
            metricseval = trainer.evaluate()

            max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(
                eval_dataset)
            metricseval["eval_samples"] = min(max_eval_samples, len(eval_dataset))

            trainer.log_metrics("eval", metricseval)
            trainer.save_metrics("eval", metricseval)

            if not os.environ.get('USE_CODALAB', 0):
                return
            import json
            json.dump(metricseval, open("dev_stats.json", "w"))

        if not args_train.do_predict:
            pass
        else:
            logger.info("Predict")
            results = trainer.predict(predict_dataset, predict_examples)
            metricspred = results.metrics

            max_predict_samples = (
                len(predict_dataset) if data_args.max_predict_samples is None else data_args.max_predict_samples
            )
            metricspred["predict_samples"] = min(max_predict_samples, len(predict_dataset))

            trainer.log_metrics("predict", metricspred)
            trainer.save_metrics("predict", metricspred)

            if not os.environ.get('USE_CODALAB', 0):
                return
            import json
            json.dump(metricspred, open("test_stats.json", "w"))

        if not args_train.push_to_hub:
            pass
        else:
            kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "question-answering"}
            if data_args.dataset_name is None:
                pass
            else:
                kwargs["dataset_tags"] = data_args.dataset_name
                if data_args.dataset_config_name is None:
                    kwargs["dataset"] = data_args.dataset_name
                else:
                    kwargs["dataset_args"] = data_args.dataset_config_name
                    kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}"

            trainer.push_to_hub(**kwargs)
        return
    raise ValueError(
        "Only models with a quick tokenizer are compatible with this sample script. To locate the model types that satisfy this criteria,"
        " visit https://huggingface.co/transformers/index.html#supported-frameworks and look at the large table of models"
    )

if __name__ == "__main__":
    main()