evaluation script for rag models import argparse import ast import logging import os import sys import pandas as pd import torch from tqdm import tqdm from transformers import bartforconditionalgeneration ragretriever ragsequenceforgeneration ragtokenforgeneration from transformers import logging as transformerslogging sys path appendos path joinos getcwd noqa e402 isort skip from utilsrag import exactmatchscore f1score noqa e402 isort skip logger logging getloggername logging basicconfiglevellogging info transformerslogging setverbosityinfo def infermodeltypemodelnameorpath if token in modelnameorpath return ragtoken if sequence in modelnameorpath return ragsequence if bart in modelnameorpath return bart return none def metricmaxovergroundtruthsmetricfn prediction groundtruths return maxmetricfnprediction gt for gt in groundtruths def getscoresargs predspath golddatapath hypos line strip for line in openpredspath r readlines answers if args golddatamode qa data pd readcsvgolddatapath sept headernone for answerlist in data1 groundtruths ast literalevalanswerlist answers appendgroundtruths else references line strip for line in opengolddatapath r readlines answers reference for reference in references f1 em total 0 for prediction groundtruths in ziphypos answers total 1 em metricmaxovergroundtruthsexactmatchscore prediction groundtruths f1 metricmaxovergroundtruthsf1score prediction groundtruths em 100 0 em total f1 100 0 f1 total logger infoff1 f1 2f logger infofem em 2f def getprecisionatkargs predspath golddatapath k args k hypos line strip for line in openpredspath r readlines references line strip for line in opengolddatapath r readlines em total 0 for hypo reference in ziphypos references hypoprovenance sethypo splitt k refprovenance setreference splitt total 1 em lenhypoprovenance refprovenance k em 100 0 em total logger infofprecisionk em 2f def evaluatebatchretrievalargs ragmodel questions def striptitletitle if title startswith title title1 if title endswith title title 1 return title retrieverinputids ragmodel retriever questionencodertokenizer batchencodeplus questions returntensorspt paddingtrue truncationtrue inputids toargs device questionencoutputs ragmodel rag questionencoderretrieverinputids questionencpooloutput questionencoutputs0 result ragmodel retriever retrieverinputids questionencpooloutput cpu detach totorch float32 numpy prefixragmodel rag generator config prefix ndocsragmodel config ndocs returntensorspt alldocs ragmodel retriever index getdocdictsresult docids provenancestrings for docs in alldocs provenance striptitletitle for title in docstitle provenancestrings appendt joinprovenance return provenancestrings def evaluatebatche2eargs ragmodel questions with torch nograd inputsdict ragmodel retriever questionencodertokenizer batchencodeplus questions returntensorspt paddingtrue truncationtrue inputids inputsdict inputids toargs device attentionmask inputsdict attentionmask toargs device outputs ragmodel generate ragmodel overwrites generate inputids attentionmaskattentionmask numbeamsargs numbeams minlengthargs minlength maxlengthargs maxlength earlystoppingfalse numreturnsequences1 badwordsids0 0 bart likes to repeat bos tokens dont allow it to generate more than one answers ragmodel retriever generatortokenizer batchdecodeoutputs skipspecialtokenstrue if args printpredictions for q a in zipquestions answers logger infoq a formatq a return answers def getargs parser argparse argumentparser parser addargument modeltype choicesragsequence ragtoken bart typestr help rag model type ragsequence ragtoken or bart if none specified the type is inferred from the modelnameorpath parser addargument indexname defaultnone choicesexact compressed legacy typestr helprag model retriever type parser addargument indexpath defaultnone typestr helppath to the retrieval index parser addargumentndocs default5 typeint helpnumber of retrieved docs parser addargument modelnameorpath defaultnone typestr requiredtrue helppath to pretrained checkpoints or model identifier from huggingface comodels parser addargument evalmode choicese2e retrieval defaulte2e typestr help evaluation mode e2e calculates exact match and f1 of the downstream task retrieval calculates precisionk parser addargumentk default1 typeint helpk for the precisionk calculation parser addargument evaluationset defaultnone typestr requiredtrue helppath to a file containing evaluation samples parser addargument golddatapath defaultnone typestr requiredtrue helppath to a tabseparated file with gold samples parser addargument golddatamode defaultqa typestr choicesqa ans help format of the gold data file qa a single line in the following format question tab answerlist ans a single line of the gold file contains the expected answer string parser addargument predictionspath typestr defaultpredictions txt helpname of the predictions file to be stored in the checkpoints directory parser addargument evalallcheckpoints actionstoretrue helpevaluate all checkpoints starting with the same prefix as modelname ending and ending with step number parser addargument evalbatchsize default8 typeint helpbatch size per gpucpu for evaluation parser addargument recalculate helprecalculate predictions even if the prediction file exists actionstoretrue parser addargument numbeams default4 typeint helpnumber of beams to be used when generating answers parser addargumentminlength default1 typeint helpmin length of the generated answers parser addargumentmaxlength default50 typeint helpmax length of the generated answers parser addargument printpredictions actionstoretrue helpif true prints predictions while evaluating parser addargument printdocs actionstoretrue helpif true prints docs retried while generating args parser parseargs args device torch devicecuda if torch cuda isavailable else cpu return args def mainargs modelkwargs if args modeltype is none args modeltype infermodeltypeargs modelnameorpath assert args modeltype is not none if args modeltype startswithrag modelclass ragtokenforgeneration if args modeltype ragtoken else ragsequenceforgeneration modelkwargsndocs args ndocs if args indexname is not none modelkwargsindexname args indexname if args indexpath is not none modelkwargsindexpath args indexpath else modelclass bartforconditionalgeneration checkpoints f path for f in os scandirargs modelnameorpath if f isdir if args evalallcheckpoints else args modelnameorpath logger infoevaluate the following checkpoints s checkpoints scorefn getscores if args evalmode e2e else getprecisionatk evaluatebatchfn evaluatebatche2e if args evalmode e2e else evaluatebatchretrieval for checkpoint in checkpoints if os path existsargs predictionspath and not args recalculate logger infocalculating metrics based on an existing predictions file formatargs predictionspath scorefnargs args predictionspath args golddatapath continue logger info running evaluation for formatcheckpoint logger info batch size d args evalbatchsize logger info predictions will be stored under formatargs predictionspath if args modeltype startswithrag retriever ragretriever frompretrainedcheckpoint modelkwargs model modelclass frompretrainedcheckpoint retrieverretriever modelkwargs model retriever initretrieval else model modelclass frompretrainedcheckpoint modelkwargs model toargs device with openargs evaluationset r as evalfile openargs predictionspath w as predsfile questions for line in tqdmevalfile questions appendline strip if lenquestions args evalbatchsize answers evaluatebatchfnargs model questions predsfile writen joinanswers n predsfile flush questions if lenquestions 0 answers evaluatebatchfnargs model questions predsfile writen joinanswers predsfile flush scorefnargs args predictionspath args golddatapath if name main args getargs mainargs noqa e402 isort skip noqa e402 isort skip rag_model overwrites generate bart likes to repeat bos tokens dont allow it to generate more than one
import argparse
import ast
import logging
import os
import sys

import pandas as pd
import torch
from tqdm import tqdm

from transformers import BartForConditionalGeneration, RagRetriever, RagSequenceForGeneration, RagTokenForGeneration
from transformers import logging as transformers_logging


sys.path.append(os.path.join(os.getcwd()))  
from utils_rag import exact_match_score, f1_score  


logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

transformers_logging.set_verbosity_info()


def infer_model_type(model_name_or_path):
    if "token" in model_name_or_path:
        return "rag_token"
    if "sequence" in model_name_or_path:
        return "rag_sequence"
    if "bart" in model_name_or_path:
        return "bart"
    return None


def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
    return max(metric_fn(prediction, gt) for gt in ground_truths)


def get_scores(args, preds_path, gold_data_path):
    hypos = [line.strip() for line in open(preds_path, "r").readlines()]
    answers = []

    if args.gold_data_mode == "qa":
        data = pd.read_csv(gold_data_path, sep="\t", header=None)
        for answer_list in data[1]:
            ground_truths = ast.literal_eval(answer_list)
            answers.append(ground_truths)
    else:
        references = [line.strip() for line in open(gold_data_path, "r").readlines()]
        answers = [[reference] for reference in references]

    f1 = em = total = 0
    for prediction, ground_truths in zip(hypos, answers):
        total += 1
        em += metric_max_over_ground_truths(exact_match_score, prediction, ground_truths)
        f1 += metric_max_over_ground_truths(f1_score, prediction, ground_truths)

    em = 100.0 * em / total
    f1 = 100.0 * f1 / total

    logger.info(f"F1: {f1:.2f}")
    logger.info(f"EM: {em:.2f}")


def get_precision_at_k(args, preds_path, gold_data_path):
    k = args.k
    hypos = [line.strip() for line in open(preds_path, "r").readlines()]
    references = [line.strip() for line in open(gold_data_path, "r").readlines()]

    em = total = 0
    for hypo, reference in zip(hypos, references):
        hypo_provenance = set(hypo.split("\t")[:k])
        ref_provenance = set(reference.split("\t"))
        total += 1
        em += len(hypo_provenance & ref_provenance) / k

    em = 100.0 * em / total
    logger.info(f"Precision@{k}: {em: .2f}")


def evaluate_batch_retrieval(args, rag_model, questions):
    def strip_title(title):
        if title.startswith('"'):
            title = title[1:]
        if title.endswith('"'):
            title = title[:-1]
        return title

    retriever_input_ids = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus(
        questions,
        return_tensors="pt",
        padding=True,
        truncation=True,
    )["input_ids"].to(args.device)

    question_enc_outputs = rag_model.rag.question_encoder(retriever_input_ids)
    question_enc_pool_output = question_enc_outputs[0]

    result = rag_model.retriever(
        retriever_input_ids,
        question_enc_pool_output.cpu().detach().to(torch.float32).numpy(),
        prefix=rag_model.rag.generator.config.prefix,
        n_docs=rag_model.config.n_docs,
        return_tensors="pt",
    )
    all_docs = rag_model.retriever.index.get_doc_dicts(result.doc_ids)
    provenance_strings = []
    for docs in all_docs:
        provenance = [strip_title(title) for title in docs["title"]]
        provenance_strings.append("\t".join(provenance))
    return provenance_strings


def evaluate_batch_e2e(args, rag_model, questions):
    with torch.no_grad():
        inputs_dict = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus(
            questions, return_tensors="pt", padding=True, truncation=True
        )

        input_ids = inputs_dict.input_ids.to(args.device)
        attention_mask = inputs_dict.attention_mask.to(args.device)
        outputs = rag_model.generate(  
            input_ids,
            attention_mask=attention_mask,
            num_beams=args.num_beams,
            min_length=args.min_length,
            max_length=args.max_length,
            early_stopping=False,
            num_return_sequences=1,
            bad_words_ids=[[0, 0]],  
        )
        answers = rag_model.retriever.generator_tokenizer.batch_decode(outputs, skip_special_tokens=True)

        if args.print_predictions:
            for q, a in zip(questions, answers):
                logger.info("Q: {} - A: {}".format(q, a))

        return answers


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_type",
        choices=["rag_sequence", "rag_token", "bart"],
        type=str,
        help=(
            "RAG model type: rag_sequence, rag_token or bart, if none specified, the type is inferred from the"
            " model_name_or_path"
        ),
    )
    parser.add_argument(
        "--index_name",
        default=None,
        choices=["exact", "compressed", "legacy"],
        type=str,
        help="RAG model retriever type",
    )
    parser.add_argument(
        "--index_path",
        default=None,
        type=str,
        help="Path to the retrieval index",
    )
    parser.add_argument("--n_docs", default=5, type=int, help="Number of retrieved docs")
    parser.add_argument(
        "--model_name_or_path",
        default=None,
        type=str,
        required=True,
        help="Path to pretrained checkpoints or model identifier from huggingface.co/models",
    )
    parser.add_argument(
        "--eval_mode",
        choices=["e2e", "retrieval"],
        default="e2e",
        type=str,
        help=(
            "Evaluation mode, e2e calculates exact match and F1 of the downstream task, retrieval calculates"
            " precision@k."
        ),
    )
    parser.add_argument("--k", default=1, type=int, help="k for the precision@k calculation")
    parser.add_argument(
        "--evaluation_set",
        default=None,
        type=str,
        required=True,
        help="Path to a file containing evaluation samples",
    )
    parser.add_argument(
        "--gold_data_path",
        default=None,
        type=str,
        required=True,
        help="Path to a tab-separated file with gold samples",
    )
    parser.add_argument(
        "--gold_data_mode",
        default="qa",
        type=str,
        choices=["qa", "ans"],
        help=(
            "Format of the gold data file"
            "qa - a single line in the following format: question [tab] answer_list"
            "ans - a single line of the gold file contains the expected answer string"
        ),
    )
    parser.add_argument(
        "--predictions_path",
        type=str,
        default="predictions.txt",
        help="Name of the predictions file, to be stored in the checkpoints directory",
    )
    parser.add_argument(
        "--eval_all_checkpoints",
        action="store_true",
        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number",
    )
    parser.add_argument(
        "--eval_batch_size",
        default=8,
        type=int,
        help="Batch size per GPU/CPU for evaluation.",
    )
    parser.add_argument(
        "--recalculate",
        help="Recalculate predictions even if the prediction file exists",
        action="store_true",
    )
    parser.add_argument(
        "--num_beams",
        default=4,
        type=int,
        help="Number of beams to be used when generating answers",
    )
    parser.add_argument("--min_length", default=1, type=int, help="Min length of the generated answers")
    parser.add_argument("--max_length", default=50, type=int, help="Max length of the generated answers")

    parser.add_argument(
        "--print_predictions",
        action="store_true",
        help="If True, prints predictions while evaluating.",
    )
    parser.add_argument(
        "--print_docs",
        action="store_true",
        help="If True, prints docs retried while generating.",
    )
    args = parser.parse_args()
    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return args


def main(args):
    model_kwargs = {}
    if args.model_type is None:
        args.model_type = infer_model_type(args.model_name_or_path)
        assert args.model_type is not None
    if args.model_type.startswith("rag"):
        model_class = RagTokenForGeneration if args.model_type == "rag_token" else RagSequenceForGeneration
        model_kwargs["n_docs"] = args.n_docs
        if args.index_name is not None:
            model_kwargs["index_name"] = args.index_name
        if args.index_path is not None:
            model_kwargs["index_path"] = args.index_path
    else:
        model_class = BartForConditionalGeneration

    checkpoints = (
        [f.path for f in os.scandir(args.model_name_or_path) if f.is_dir()]
        if args.eval_all_checkpoints
        else [args.model_name_or_path]
    )

    logger.info("Evaluate the following checkpoints: %s", checkpoints)

    score_fn = get_scores if args.eval_mode == "e2e" else get_precision_at_k
    evaluate_batch_fn = evaluate_batch_e2e if args.eval_mode == "e2e" else evaluate_batch_retrieval

    for checkpoint in checkpoints:
        if os.path.exists(args.predictions_path) and (not args.recalculate):
            logger.info("Calculating metrics based on an existing predictions file: {}".format(args.predictions_path))
            score_fn(args, args.predictions_path, args.gold_data_path)
            continue

        logger.info("***** Running evaluation for {} *****".format(checkpoint))
        logger.info("  Batch size = %d", args.eval_batch_size)
        logger.info("  Predictions will be stored under {}".format(args.predictions_path))

        if args.model_type.startswith("rag"):
            retriever = RagRetriever.from_pretrained(checkpoint, **model_kwargs)
            model = model_class.from_pretrained(checkpoint, retriever=retriever, **model_kwargs)
            model.retriever.init_retrieval()
        else:
            model = model_class.from_pretrained(checkpoint, **model_kwargs)
        model.to(args.device)

        with open(args.evaluation_set, "r") as eval_file, open(args.predictions_path, "w") as preds_file:
            questions = []
            for line in tqdm(eval_file):
                questions.append(line.strip())
                if len(questions) == args.eval_batch_size:
                    answers = evaluate_batch_fn(args, model, questions)
                    preds_file.write("\n".join(answers) + "\n")
                    preds_file.flush()
                    questions = []
            if len(questions) > 0:
                answers = evaluate_batch_fn(args, model, questions)
                preds_file.write("\n".join(answers))
                preds_file.flush()

            score_fn(args, args.predictions_path, args.gold_data_path)


if __name__ == "__main__":
    args = get_args()
    main(args)