chendl/compositional_test

1""" Evaluation script for RAG models."""2 3import argparse4import ast5import logging6import os7import sys8 9import pandas as pd10import torch11from tqdm import tqdm12 13from transformers import BartForConditionalGeneration, RagRetriever, RagSequenceForGeneration, RagTokenForGeneration14from transformers import logging as transformers_logging15 16 17sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # isort:skip18from utils_rag import exact_match_score, f1_score  # noqa: E402 # isort:skip19 20 21logger = logging.getLogger(__name__)22logging.basicConfig(level=logging.INFO)23 24transformers_logging.set_verbosity_info()25 26 27def infer_model_type(model_name_or_path):28    if "token" in model_name_or_path:29        return "rag_token"30    if "sequence" in model_name_or_path:31        return "rag_sequence"32    if "bart" in model_name_or_path:33        return "bart"34    return None35 36 37def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):38    return max(metric_fn(prediction, gt) for gt in ground_truths)39 40 41def get_scores(args, preds_path, gold_data_path):42    hypos = [line.strip() for line in open(preds_path, "r").readlines()]43    answers = []44 45    if args.gold_data_mode == "qa":46        data = pd.read_csv(gold_data_path, sep="\t", header=None)47        for answer_list in data[1]:48            ground_truths = ast.literal_eval(answer_list)49            answers.append(ground_truths)50    else:51        references = [line.strip() for line in open(gold_data_path, "r").readlines()]52        answers = [[reference] for reference in references]53 54    f1 = em = total = 055    for prediction, ground_truths in zip(hypos, answers):56        total += 157        em += metric_max_over_ground_truths(exact_match_score, prediction, ground_truths)58        f1 += metric_max_over_ground_truths(f1_score, prediction, ground_truths)59 60    em = 100.0 * em / total61    f1 = 100.0 * f1 / total62 63    logger.info(f"F1: {f1:.2f}")64    logger.info(f"EM: {em:.2f}")65 66 67def get_precision_at_k(args, preds_path, gold_data_path):68    k = args.k69    hypos = [line.strip() for line in open(preds_path, "r").readlines()]70    references = [line.strip() for line in open(gold_data_path, "r").readlines()]71 72    em = total = 073    for hypo, reference in zip(hypos, references):74        hypo_provenance = set(hypo.split("\t")[:k])75        ref_provenance = set(reference.split("\t"))76        total += 177        em += len(hypo_provenance & ref_provenance) / k78 79    em = 100.0 * em / total80    logger.info(f"Precision@{k}: {em: .2f}")81 82 83def evaluate_batch_retrieval(args, rag_model, questions):84    def strip_title(title):85        if title.startswith('"'):86            title = title[1:]87        if title.endswith('"'):88            title = title[:-1]89        return title90 91    retriever_input_ids = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus(92        questions,93        return_tensors="pt",94        padding=True,95        truncation=True,96    )["input_ids"].to(args.device)97 98    question_enc_outputs = rag_model.rag.question_encoder(retriever_input_ids)99    question_enc_pool_output = question_enc_outputs[0]100 101    result = rag_model.retriever(102        retriever_input_ids,103        question_enc_pool_output.cpu().detach().to(torch.float32).numpy(),104        prefix=rag_model.rag.generator.config.prefix,105        n_docs=rag_model.config.n_docs,106        return_tensors="pt",107    )108    all_docs = rag_model.retriever.index.get_doc_dicts(result.doc_ids)109    provenance_strings = []110    for docs in all_docs:111        provenance = [strip_title(title) for title in docs["title"]]112        provenance_strings.append("\t".join(provenance))113    return provenance_strings114 115 116def evaluate_batch_e2e(args, rag_model, questions):117    with torch.no_grad():118        inputs_dict = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus(119            questions, return_tensors="pt", padding=True, truncation=True120        )121 122        input_ids = inputs_dict.input_ids.to(args.device)123        attention_mask = inputs_dict.attention_mask.to(args.device)124        outputs = rag_model.generate(  # rag_model overwrites generate125            input_ids,126            attention_mask=attention_mask,127            num_beams=args.num_beams,128            min_length=args.min_length,129            max_length=args.max_length,130            early_stopping=False,131            num_return_sequences=1,132            bad_words_ids=[[0, 0]],  # BART likes to repeat BOS tokens, dont allow it to generate more than one133        )134        answers = rag_model.retriever.generator_tokenizer.batch_decode(outputs, skip_special_tokens=True)135 136        if args.print_predictions:137            for q, a in zip(questions, answers):138                logger.info("Q: {} - A: {}".format(q, a))139 140        return answers141 142 143def get_args():144    parser = argparse.ArgumentParser()145    parser.add_argument(146        "--model_type",147        choices=["rag_sequence", "rag_token", "bart"],148        type=str,149        help=(150            "RAG model type: rag_sequence, rag_token or bart, if none specified, the type is inferred from the"151            " model_name_or_path"152        ),153    )154    parser.add_argument(155        "--index_name",156        default=None,157        choices=["exact", "compressed", "legacy"],158        type=str,159        help="RAG model retriever type",160    )161    parser.add_argument(162        "--index_path",163        default=None,164        type=str,165        help="Path to the retrieval index",166    )167    parser.add_argument("--n_docs", default=5, type=int, help="Number of retrieved docs")168    parser.add_argument(169        "--model_name_or_path",170        default=None,171        type=str,172        required=True,173        help="Path to pretrained checkpoints or model identifier from huggingface.co/models",174    )175    parser.add_argument(176        "--eval_mode",177        choices=["e2e", "retrieval"],178        default="e2e",179        type=str,180        help=(181            "Evaluation mode, e2e calculates exact match and F1 of the downstream task, retrieval calculates"182            " precision@k."183        ),184    )185    parser.add_argument("--k", default=1, type=int, help="k for the precision@k calculation")186    parser.add_argument(187        "--evaluation_set",188        default=None,189        type=str,190        required=True,191        help="Path to a file containing evaluation samples",192    )193    parser.add_argument(194        "--gold_data_path",195        default=None,196        type=str,197        required=True,198        help="Path to a tab-separated file with gold samples",199    )200    parser.add_argument(201        "--gold_data_mode",202        default="qa",203        type=str,204        choices=["qa", "ans"],205        help=(206            "Format of the gold data file"207            "qa - a single line in the following format: question [tab] answer_list"208            "ans - a single line of the gold file contains the expected answer string"209        ),210    )211    parser.add_argument(212        "--predictions_path",213        type=str,214        default="predictions.txt",215        help="Name of the predictions file, to be stored in the checkpoints directory",216    )217    parser.add_argument(218        "--eval_all_checkpoints",219        action="store_true",220        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number",221    )222    parser.add_argument(223        "--eval_batch_size",224        default=8,225        type=int,226        help="Batch size per GPU/CPU for evaluation.",227    )228    parser.add_argument(229        "--recalculate",230        help="Recalculate predictions even if the prediction file exists",231        action="store_true",232    )233    parser.add_argument(234        "--num_beams",235        default=4,236        type=int,237        help="Number of beams to be used when generating answers",238    )239    parser.add_argument("--min_length", default=1, type=int, help="Min length of the generated answers")240    parser.add_argument("--max_length", default=50, type=int, help="Max length of the generated answers")241 242    parser.add_argument(243        "--print_predictions",244        action="store_true",245        help="If True, prints predictions while evaluating.",246    )247    parser.add_argument(248        "--print_docs",249        action="store_true",250        help="If True, prints docs retried while generating.",251    )252    args = parser.parse_args()253    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")254    return args255 256 257def main(args):258    model_kwargs = {}259    if args.model_type is None:260        args.model_type = infer_model_type(args.model_name_or_path)261        assert args.model_type is not None262    if args.model_type.startswith("rag"):263        model_class = RagTokenForGeneration if args.model_type == "rag_token" else RagSequenceForGeneration264        model_kwargs["n_docs"] = args.n_docs265        if args.index_name is not None:266            model_kwargs["index_name"] = args.index_name267        if args.index_path is not None:268            model_kwargs["index_path"] = args.index_path269    else:270        model_class = BartForConditionalGeneration271 272    checkpoints = (273        [f.path for f in os.scandir(args.model_name_or_path) if f.is_dir()]274        if args.eval_all_checkpoints275        else [args.model_name_or_path]276    )277 278    logger.info("Evaluate the following checkpoints: %s", checkpoints)279 280    score_fn = get_scores if args.eval_mode == "e2e" else get_precision_at_k281    evaluate_batch_fn = evaluate_batch_e2e if args.eval_mode == "e2e" else evaluate_batch_retrieval282 283    for checkpoint in checkpoints:284        if os.path.exists(args.predictions_path) and (not args.recalculate):285            logger.info("Calculating metrics based on an existing predictions file: {}".format(args.predictions_path))286            score_fn(args, args.predictions_path, args.gold_data_path)287            continue288 289        logger.info("***** Running evaluation for {} *****".format(checkpoint))290        logger.info("  Batch size = %d", args.eval_batch_size)291        logger.info("  Predictions will be stored under {}".format(args.predictions_path))292 293        if args.model_type.startswith("rag"):294            retriever = RagRetriever.from_pretrained(checkpoint, **model_kwargs)295            model = model_class.from_pretrained(checkpoint, retriever=retriever, **model_kwargs)296            model.retriever.init_retrieval()297        else:298            model = model_class.from_pretrained(checkpoint, **model_kwargs)299        model.to(args.device)300 301        with open(args.evaluation_set, "r") as eval_file, open(args.predictions_path, "w") as preds_file:302            questions = []303            for line in tqdm(eval_file):304                questions.append(line.strip())305                if len(questions) == args.eval_batch_size:306                    answers = evaluate_batch_fn(args, model, questions)307                    preds_file.write("\n".join(answers) + "\n")308                    preds_file.flush()309                    questions = []310            if len(questions) > 0:311                answers = evaluate_batch_fn(args, model, questions)312                preds_file.write("\n".join(answers))313                preds_file.flush()314 315            score_fn(args, args.predictions_path, args.gold_data_path)316 317 318if __name__ == "__main__":319    args = get_args()320    main(args)321