[Data + code] ExpertQA : Expert-Curated Questions and Attributed Answers
"""Utilities for reading and writing examples."""

import dataclasses
import typing
import os
import sys

sys.path.append(os.path.dirname(os.path.realpath(__file__)))

from data_utils import jsonl_utils


@dataclasses.dataclass(frozen=False)
class Claim:
    """Represents a claim."""
    # Original claim
    claim_string: str
    # Evidence for claim
    evidence: typing.Optional[typing.Sequence[str]] = None
    # Supported (by annotator)
    support: typing.Optional[int] = None
    # Reason for missing support (by annotator)
    reason_missing_support: typing.Optional[str] = None
    # Informativeness of claim for the question (by annotator)
    informativeness: typing.Optional[int] = None
    # Worthiness of citing claim (by annotator)
    worthiness: typing.Optional[int] = None
    # Factual correctness of claim (by annotator)
    correctness: typing.Optional[int] = None
    # Reliability of source (by annotator)
    reliability: typing.Optional[int] = None
    # Revised claim (by annotator)
    revised_claim: typing.Optional[str] = None
    # Revised evidence (by annotator)
    revised_evidence: typing.Optional[typing.Sequence[str]] = None
    # Atomic claims for Fact score estimation
    atomic_claims: typing.Optional[typing.Sequence[str]] = None
    # Atomic claim-evidences for Fact score estimation
    atomic_evidences: typing.Optional[typing.Sequence[str]] = None
    # Fact score for each claim
    fact_score: typing.Optional[float] = None
    # Autoais label for the original claim and original evidence
    autoais_label: typing.Optional[str] = None


@dataclasses.dataclass(frozen=False)
class Answer:
    """Represents an answer."""
    # Original answer returned by system
    answer_string: str
    # Attribution (not linked to specific claims)
    attribution: typing.Optional[typing.Sequence[str]] = None
    # List of claims for the answer
    claims: typing.Optional[typing.Sequence[Claim]] = None
    # Revised answer (by annotator)
    revised_answer_string: typing.Optional[str] = None
    # Usefulness of original answer (by annotator)
    usefulness: typing.Optional[int] = None
    # Time taken for annotating this answer (from annotation)
    annotation_time: typing.Optional[float] = None
    # Prolific ID (TODO: anonymize before release) or email ID of annotator who annotated answer
    annotator_id: typing.Optional[str] = None


@dataclasses.dataclass(frozen=False)
class ExampleMetadata:
    """Optional metadata."""
    # The type of the question.
    question_type: str = None
    # The field of the expert who wrote the the question.
    field: str = None
    # The specific area in the field (if specified)
    specific_field: typing.Optional[str] = None


@dataclasses.dataclass(frozen=False)
class Example:
    """Represents a question."""
    question: str
    # Prolific ID (TODO: anonymize before release) or email ID of annotator who wrote the question
    annotator_id: str
    # Dict of model name to answer string
    answers: typing.Optional[typing.DefaultDict[str, Answer]] = None
    # Optional metadata.
    metadata: typing.Optional[ExampleMetadata] = None


def read_examples(filepath):
    examples_json = jsonl_utils.read_jsonl(filepath)
    examples = [Example(**example) for example in examples_json]
    for example in examples:
        example.metadata = ExampleMetadata(**example.metadata)
    return examples


def write_examples(filepath, examples, append=False):
    examples_json = [dataclasses.asdict(example) for example in examples]
    jsonl_utils.write_jsonl(filepath, examples_json, append)

data_utils/example_utils.py
chaitanyamalaviya/ExpertQA
import json


def read_jsonl(filepath, limit=None, verbose=False):
    """Read jsonl file to a List of Dicts."""
    data = []
    with open(filepath, "r") as jsonl_file:
      for idx, line in enumerate(jsonl_file):
          if limit is not None and idx >= limit:
              break
          if verbose and idx % 100 == 0:
              # Print the index every 100 lines.
              print("Processing line %s." % idx)
          try:
              data.append(json.loads(line))
          except json.JSONDecodeError as e:
              print("Failed to parse line: `%s`" % line)
              raise e
    print("Loaded %s lines from %s." % (len(data), filepath))
    return data


def write_jsonl(filepath, rows, append=False):
    """Write a List of Dicts to jsonl file."""
    if append:
        jsonl_file = open(filepath, "a")
    else:
        jsonl_file = open(filepath, "w")
    for row in rows:
        line = "%s\n" % json.dumps(row)
        jsonl_file.write(line)
    print("Wrote %s lines to %s." % (len(rows), filepath))

data_utils/jsonl_utils.py
chaitanyamalaviya/ExpertQA
from absl import app
from absl import flags
from qafacteval import QAFactEval
import os
import sys
from evaluate import load

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from data_utils import example_utils, jsonl_utils

FLAGS = flags.FLAGS

flags.DEFINE_string('prediction_file', 'saved_models/domain_expertqa_flant5_xxl/test_eval_outputs.json', 'Prediction file.')
flags.DEFINE_string('reference_file', 'data/expertqa/domain_lfqa_test.json', 'Reference file.')


def calculate_rouge(predictions, references):
    rouge = load("rouge")
    results = rouge.compute(predictions=predictions, references=references)
    print(results)

def calculate_qafacteval_score(predictions, references):
    
    kwargs = {"cuda_device": 0, "use_lerc_quip": True, \
            "verbose": True, "generation_batch_size": 32, \
            "answering_batch_size": 32, "lerc_batch_size": 8}

    model_folder = "eval/qafacteval_models" # path to models downloaded with download_models.sh
    metric = QAFactEval(
        lerc_quip_path=f"{model_folder}/quip-512-mocha",
        generation_model_path=f"{model_folder}/generation/model.tar.gz",
        answering_model_dir=f"{model_folder}/answering",
        lerc_model_path=f"{model_folder}/lerc/model.tar.gz",
        lerc_pretrained_model_path=f"{model_folder}/lerc/pretraining.tar.gz",
        **kwargs
    )

    results = metric.score_batch_qafacteval(references, [predictions], return_qa_pairs=True)
    score = results[0][0]['qa-eval']['lerc_quip']
    print(score)


def main(unused_argv):
    # Example usage
    predictions = [ex["output"][0] for ex in jsonl_utils.read_jsonl(FLAGS.prediction_file)]
    references = [ex["answer"] for ex in jsonl_utils.read_jsonl(FLAGS.reference_file)]
    
    calculate_rouge(predictions, references)
    calculate_qafacteval_score(predictions, references)


if __name__ == "__main__":
    app.run(main)
eval/nlg_eval.py
chaitanyamalaviya/ExpertQA
"""Evaluation of Attributed QA systems."""
import dataclasses
import json
import os
import re
import sys

from absl import app
from absl import flags
from absl import logging
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
import numpy as np
from tqdm import tqdm
from transformers import T5ForConditionalGeneration, T5Config
from transformers import T5Tokenizer
import torch

from nltk import sent_tokenize

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils


AUTOAIS = "google/t5_xxl_true_nli_mixture"

PASSAGE_FORMAT = re.compile("« ([^»]*) » « ([^»]*) » (.*)")

# File locations.
flags.DEFINE_string("input_file", "", "Path to system predictions.")
flags.DEFINE_string("ais_output_file", "", "Path for file to write AIS outputs.")
flags.DEFINE_string("t5_ckpt_path", None, "Path to finetuned T5-xxl TRUE checkpoint.")
FLAGS = flags.FLAGS


def format_passage_for_autoais(passage):
    """Produce the NLI format for a passage.

    Args:
        passage: A passage from the Wikipedia scrape.

    Returns:
        a formatted string, e.g.

        Luke Cage (season 2), Release. The second season of Luke Cage was released
        on June 22, 2018, on the streaming service Netflix worldwide, in Ultra HD
        4K and high dynamic range.
    """
    m = PASSAGE_FORMAT.match(passage)
    if not m:
        return passage

    headings = m.group(2)
    passage = m.group(3)
    return f"{headings}. {passage}"


def format_example_for_autoais(evidence, claim):
    return "premise: {} hypothesis: {}".format(evidence, claim)


def get_entail_label_ids(tokenizer):
    pos_id = tokenizer.convert_tokens_to_ids("1")
    neg_id = tokenizer.convert_tokens_to_ids("0")
    return pos_id, neg_id


def _softmax(x):
    return np.exp(x) / np.sum(np.exp(x), axis=0)


def _autoais_predict(evidence,
                     claim,
                     model,
                     tokenizer,
                     entail_token_id,
                     non_entail_token_id,
                     cuda: bool = True):
    input_text = format_example_for_autoais(evidence, claim)
    input_ids = tokenizer(input_text,
                          return_tensors="pt",
                          padding="max_length",
                          max_length=512,
                          truncation=True).input_ids

    if cuda:
        input_ids = input_ids.cuda()

    outputs = model.generate(input_ids, return_dict_in_generate=True, output_scores=True)
    result = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
    inference = "Y" if result == "1" else "N"

    # Get softmax score between "1" and "0"
    tok_scores = outputs.scores[0][0].detach().cpu()
    pos_score = tok_scores[entail_token_id].numpy()
    neg_score = tok_scores[non_entail_token_id].numpy()
    label_scores = _softmax([pos_score, neg_score])

    return inference, label_scores


def stretch_nli_rerank(example,
                       tokenizer,
                       model,
                       entail_token_id: int,
                       non_entail_token_id: int,
                       k: int = 2
                       ):
    """Slice evidence into smaller slices and re-rank the slices based on model confidence.

    See https://arxiv.org/abs/2204.07447 for details.
    Since the AutoAIS model doesn't distinguish between contradiction and netural,

    """

    # Slice evidence into smaller chunks
    claim_text = example["claim"]
    evidence_text = example["evidence"]
    sents = sent_tokenize(evidence_text)

    per_sent_score = []
    for sent in sents:
        input_text = format_example_for_autoais(sent, claim_text)
        label, score = _autoais_predict(evidence=input_text,
                                        claim=claim_text,
                                        model=model,
                                        tokenizer=tokenizer,
                                        entail_token_id=entail_token_id,
                                        non_entail_token_id=non_entail_token_id)
        per_sent_score.append(score[1]) # order by neutral prob, so that argsort by ascending order will give us
                                        # evidence sentences from most-to-least entailed


    # Take top k entailment sentences + top k neutral sentences, run through the autoais model again
    high_to_low = np.argsort(per_sent_score)
    reranked_indices = np.concatenate((high_to_low[:k], high_to_low[-k:]), axis=None)
    reranked_input_sents = [sents[idx] for idx in reranked_indices]

    reranked_input = " ".join(reranked_input_sents)
    label, score = _autoais_predict(evidence=reranked_input,
                                    claim=claim_text,
                                    model=model,
                                    tokenizer=tokenizer,
                                    entail_token_id=entail_token_id,
                                    non_entail_token_id=non_entail_token_id)

    return label


def infer_autoais(example,
                  tokenizer,
                  model,
                  max_len:int = 512):
    """Runs inference for assessing AIS between a premise and hypothesis.

    Args:
        example: Dict with the example data.
        tokenizer: A huggingface tokenizer object.
        model: A huggingface model object.

    Returns:
        A string representing the model prediction.
    """
    input_text = format_example_for_autoais(example["evidence"], example["claim"])
    pos_id, neg_id = get_entail_label_ids(tokenizer)

    # Try tokenize the input text, see if it exceeds the token limit
    tokens = tokenizer.encode(input_text)
    if len(tokens) > max_len:
        inference = stretch_nli_rerank(
            example=example,
            tokenizer=tokenizer,
            model=model,
            entail_token_id=pos_id,
            non_entail_token_id=neg_id,
            k=2
        )
    else:
        inference, _ = _autoais_predict(
            evidence=example["evidence"],
            claim=example["claim"],
            model=model,
            tokenizer=tokenizer,
            entail_token_id=pos_id,
            non_entail_token_id=neg_id
        )

    return inference


def remove_urls(input_string):
    # Regular expression to match URLs
    url_pattern = r'https?://\S+|www\.\S+'
    # Replace URLs with an empty string
    cleaned_string = re.sub(url_pattern, '', input_string)
    return cleaned_string


def cleanup_claim_and_evidence(raw_text):
    """
    Remove square brackets and urls in claim and evidence
    """
    # remove square brackets in the string
    cleaned_str = re.sub(r'\[\d\]\s?', '', raw_text)

    # remove urls in evidence
    parts = cleaned_str.split("\n\n")
    if parts[0].startswith("http"):
        cleaned_str = " ".join(parts[1:])
    else:
        cleaned_str = " ".join(parts)

    return cleaned_str


def score_predictions(input_data,
                      output_path,
                      t5_ckpt_path: str = None):
    """Scores model predictions using AutoAIS.

    Args:
        input_data: list of input example dicts with claim and evidence keys

    Returns:
        input data examples along with autoais_prediction for each example
    """
    hf_tokenizer = T5Tokenizer.from_pretrained(AUTOAIS, cache_dir="/nlp/data/huggingface_cache/")

    if t5_ckpt_path is None:
        hf_model = T5ForConditionalGeneration.from_pretrained(AUTOAIS, device_map='auto',
                                                              cache_dir="/nlp/data/huggingface_cache/")
    else:
        cfg = T5Config.from_pretrained(AUTOAIS)
        with init_empty_weights():
            hf_model = T5ForConditionalGeneration._from_config(cfg)

        hf_model = load_checkpoint_and_dispatch(hf_model, t5_ckpt_path, device_map='auto')

    with open(output_path, "w") as fout:
        for i, ex in enumerate(tqdm(input_data)):
            for system, answer in ex.answers.items():
                for c in answer["claims"]:
                    claim_string = c["claim_string"]
                    claim_string = cleanup_claim_and_evidence(claim_string)
                    if not c["evidence"] or len(c["evidence"]) == 0:
                        print(f"empty evidence for {claim_string}!")
                        c["autoais_label"] = "N"
                        continue
                    evidence = [cleanup_claim_and_evidence(e) for e in c["evidence"]]
                    evidence_joined = " ".join(evidence)

                    input_ex = {
                        "claim": claim_string,
                        "evidence": evidence_joined
                    }
                    autoais_label = infer_autoais(input_ex, hf_tokenizer, hf_model)
                    c["autoais_label"] = autoais_label

            ex_json = dataclasses.asdict(ex)
            ex_str = json.dumps(ex_json)
            fout.write(ex_str)
            fout.write("\n")


def main(unused_argv):
    input_data = example_utils.read_examples(FLAGS.input_file)

    logging.info("Scoring and writing predictions...")
    score_predictions(input_data, FLAGS.ais_output_file, FLAGS.t5_ckpt_path)


if __name__ == "__main__":
  app.run(main)
modeling/auto_attribution/autoais.py
chaitanyamalaviya/ExpertQA
'''Given a file where autoais score has already been added to each claim, compute dataset-wise autoais score. '''
import numpy as np
import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils


def did_abstain(answer):
    if "The context does not" in answer["answer_string"] or \
        "I don't know" in answer["answer_string"] or \
        "The context provided does not" in answer["answer_string"] or \
        "The context provided doesn't" in answer["answer_string"] or \
        "The context doesn't" in answer["answer_string"] or \
        "not sure if I understand" in answer["answer_string"] or \
        "I'm sorry, but I can't" in answer["answer_string"] or \
        "The given context does not" in answer["answer_string"] or \
        "I couldn't find any relevant information" in answer["answer_string"]:
        return True
    else:
        return False

def print_autoais_for_system(predictions):
    # Compute macro- (question-level) and micro- (claim-level) averaged autoais score over all predictions for a system
    all_autoais_scores = {}
    all_autoais_scores_without_abstain = {}
    all_revised_autoais_scores = {}
    abstain_questions_count = {}
    empty_evidence_count = {}
    total_claim_count = {}
    valid_claim_count = {}
    for ex in predictions:
        for system, answer in ex.answers.items():
            if system not in all_autoais_scores:
                all_autoais_scores[system] = []
                all_revised_autoais_scores[system] = []
                empty_evidence_count[system] = 0
                total_claim_count[system] = 0
                valid_claim_count[system] = 0
                all_autoais_scores_without_abstain[system] = []
                abstain_questions_count[system] = 0

            per_answer_autoais = []
            is_abstain = did_abstain(answer)
            for c in answer["claims"]:
                total_claim_count[system] += 1

                # if c["support"] != "Complete":
                #     continue
                # Skip the claim if there's no evidence, or if the evidence is too short
                # TODO(sihaoc): Change this if this is not what we want
                if not c["evidence"] or len(c["evidence"]) == 0:
                    empty_evidence_count[system] += 1
                    continue

                evidence_joined = "|".join(c["evidence"])
                if "http" in evidence_joined and system == "gpt4":  # for GPT4, we discount evidence without valid links
                    continue

                if c["support"] != "Complete" or c["worthiness"] != "Yes":
                    continue

                valid_claim_count[system] += 1

                if c["autoais_label"] == "Y":
                    per_answer_autoais.append(1)
                else:
                    per_answer_autoais.append(0)

            all_autoais_scores[system].append(per_answer_autoais)
            if not is_abstain:
                all_autoais_scores_without_abstain[system].append(per_answer_autoais)
            else:
                abstain_questions_count[system] += 1


    for system, scores in all_autoais_scores.items():
        # macro avg (avg autoais across question/answer)
        all_answer_autoais = [np.mean(ex) for ex in scores if not np.isnan(np.mean(ex))]
        # all_answer_autoais = np.nan_to_num(all_answer_autoais)  # if the answer doesn't have a claim
        #                                                         # Assign avg ais score of 0 to the answer
        macro_autoais = np.mean(all_answer_autoais)

        scores_wo_abstain = all_autoais_scores_without_abstain[system]
        all_answer_autoais_wo_abstain = [np.mean(ex) for ex in scores_wo_abstain if not np.isnan(np.mean(ex))]
        # all_answer_autoais_wo_abstain = np.nan_to_num(all_answer_autoais_wo_abstain)
        macro_autoais_wo_abstain = np.mean(all_answer_autoais_wo_abstain)

        # micro avg (avg autoais across claims)
        all_autoais_scores_flat = []
        all_autoais_scores_flat_wo_abstain = []
        for ex in scores:
            all_autoais_scores_flat += ex
        for ex in scores_wo_abstain:
            all_autoais_scores_flat_wo_abstain += ex
        micro_autoais = np.mean(all_autoais_scores_flat)
        micro_autoais_wo_abstain = np.mean(all_autoais_scores_flat_wo_abstain)

        print(f"System: {system}")
        print(f"\tMacro Avg. (avg autoais across questions): {macro_autoais:.3f}")
        print(f"\tMicro Avg. (avg autoais across claims): {micro_autoais:.3f}")
        print(f"\tNumber of claims without evidence: {empty_evidence_count[system]} / {total_claim_count[system]}")
        print(f"\tAbstain question count: {abstain_questions_count[system]} / {len(all_autoais_scores[system])}")
        print(f"\tClaims worthy + has complete support: {valid_claim_count[system]}")
        print("\tPerformance without abstained answers (i.e. excluding all 'I don't know' answers)")
        print(f"\t\tMacro Avg. (avg autoais across questions): {macro_autoais_wo_abstain:.3f}")
        print(f"\t\tMicro Avg. (avg autoais across claims): {micro_autoais_wo_abstain:.3f}")


def main(args):
    print("Reading system predictions with autoais score per claim...")
    for fin in args.prediction_files:
        input_data = example_utils.read_examples(fin)
        print_autoais_for_system(input_data)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--prediction_files", nargs='+', help="bar", type=str)
    args = parser.parse_args()

    main(args)
modeling/auto_attribution/compute_autoais_score.py
chaitanyamalaviya/ExpertQA
"""
Compute AutoAIS labels with respect to human judgements.
"""
from data_utils.jsonl_utils import read_jsonl
from sklearn.metrics import accuracy_score


def human_corr(autoais_test_file):
    examples = read_jsonl(autoais_test_file)

    preds = []
    golds = []

    per_sys_preds = {}
    per_sys_golds = {}
    for ex in examples:
        for system, answer in ex["answers"].items():
            if system not in per_sys_preds:
                per_sys_preds[system] = []
                per_sys_golds[system] = []

            for c in answer["claims"]:
                # If a claim is not worthy of citation, we skip it in the evaluation
                if c["worthiness"] == "No":
                    continue

                autoais_label = 1 if c["autoais_label"] == "Y" else 0
                gold_label = 1 if c["support"] == "Complete" else 0

                preds.append(autoais_label)
                golds.append(gold_label)
                per_sys_preds[system].append(autoais_label)
                per_sys_golds[system].append(gold_label)

    print(f"Accuracy: {accuracy_score(golds, preds)}")
    print(f"Num positives in gold = {sum(golds)} / {len(golds)}")

    for system, sys_golds in per_sys_golds.items():
        sys_preds = per_sys_preds[system]
        print(f"\t{system}")
        print(f"\tAccuracy: {accuracy_score(sys_golds, sys_preds)}")
        print(f"\tNum positives in gold = {sum(sys_golds)} / {len(sys_preds)}")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--autoais_test_set",
                        default="modeling/auto_attribution/domain_test_autoais_0shot.jsonl",
                        help="example file with per-claim autoais labels",
                        type=str)
    args = parser.parse_args()

    human_corr(args.autoais_test_set)

modeling/auto_attribution/compute_human_correlation.py
chaitanyamalaviya/ExpertQA
'''Script to convert examples for evaluating with AutoAIS'''

from absl import app
from absl import flags
from absl import logging
import os
import re
import sys
from newspaper import Article
from newspaper.article import ArticleException
import numpy as np
from tqdm import tqdm

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import jsonl_utils

FLAGS = flags.FLAGS

flags.DEFINE_string('input_file', 'data/r1_data_answers_bingchat_balanced_final_claims.jsonl', 'Input filepath.')
flags.DEFINE_string('output_file', 'data/r1_data_answers_bingchat_balanced_final_claims_autoais.json', 'Output filepath.')


def scrape_and_parse(url: str):
    """Scrape a webpage and parse it into a Document object"""
    a = Article(url)
    try:
        a.download()
        a.parse()
    except ArticleException:
        return None

    return {
        "url": url,
        "text": a.text,
    }


def main(unused_argv):
    input_data = jsonl_utils.read_jsonl(FLAGS.input_file)
    corrected_data = []
    evidence_tok_len = {}
    processed_claim_counts = 0
    for i, ex in enumerate(tqdm(input_data)):

        for answer_model_name, answer in ex["answers"].items():
            # Check if the system is bingchat or gpt4
            # Only scrape evidence for the two systems
            if answer_model_name not in ["bing_chat", "gpt4"]:
                continue

            if answer_model_name not in evidence_tok_len:
                evidence_tok_len[answer_model_name] = []

            for claim in answer["claims"]:
                if "evidence" not in claim or len(claim["evidence"]) == 0:
                    continue
                urls = [re.search("(?P<url>https?://[^\s]+)", ev).group("url") for ev in claim["evidence"]]
                evidences = [scrape_and_parse(url) for url in urls]
                if evidences is None or len(evidences) == 0:
                    continue

                evidence_text = [ev["text"] for ev in evidences if ev is not None and "text" in ev and ev["text"]]
                claim["evidence"] = evidence_text
                evidence_tok_len[answer_model_name].append(sum([len(ev.split()) for ev in evidence_text]))
                processed_claim_counts += 1

        corrected_data.append(ex)

    logging.info("Crawled evidence for %d claims...", processed_claim_counts)
    jsonl_utils.write_jsonl(FLAGS.output_file, corrected_data)

    # log the evidence token counts for each system
    for answer_model_name, counts in evidence_tok_len.items():
        logging.info("Avg. evidence token counts for %s: %f", answer_model_name, float(np.mean(counts)))


if __name__ == "__main__":
    app.run(main)
modeling/auto_attribution/convert_for_autoais.py
chaitanyamalaviya/ExpertQA
"""
Given the full r2 data, and a subset (e.g. domain val/test), extract the subset from the full r2 data.

This is due to the fact that gpt4 and bing-chat in the official train/test/val split only contains evidence urls, not
the full text. We want to extract the processed evidence text + autoais labels from the full r2 data.
"""
from data_utils.jsonl_utils import read_jsonl, write_jsonl


def extract_subset(args):
    r2_data = read_jsonl(args.r2_compiled)
    subset_data = read_jsonl(args.subset)

    r2_data_map = {ex["question"]: ex for ex in r2_data}

    all_examples = []
    for ex in subset_data:
        if ex["question"] not in r2_data_map:
            print(f"{ex['question']} not found in r2 data; This shouldn't happen!!")
            continue

        new_ex = r2_data_map[ex["question"]]
        all_examples.append(new_ex)

    write_jsonl(args.output, all_examples)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--r2_compiled",
                        default="modeling/auto_attribution/r2_compiled_out_corrected_autoais.jsonl",
                        help="full r2 compiled data",
                        type=str)
    parser.add_argument("--subset",
                        default="data/expertqa/domain_test.jsonl",
                        help="subset of data",
                        type=str)
    parser.add_argument("--output",
                        default="modeling/auto_attribution/domain_test_autoais_0shot.jsonl",
                        help="path to output",
                        type=str)
    args = parser.parse_args()

    extract_subset(args)


modeling/auto_attribution/extract_subset.py
chaitanyamalaviya/ExpertQA
"""Finetune AutoAIS model on human judgements in ExpertQA. """

import json
import os
import argparse
import numpy as np
from functools import partial
from transformers import (
    AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq,
    AutoTokenizer,
    set_seed,
)
from datasets import Dataset
import torch
import evaluate
import nltk
import numpy as np

from huggingface_hub import HfFolder
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

from data_utils.jsonl_utils import read_jsonl
from modeling.auto_attribution.autoais import cleanup_claim_and_evidence, format_example_for_autoais

nltk.download("punkt", quiet=True)

# Metric
metric = evaluate.load("rouge")
# evaluation generation args
gen_kwargs = {
    "early_stopping": True,
    "length_penalty": 2.0,
    "max_new_tokens": 50,
    "min_length": 30,
    "no_repeat_ngram_size": 3,
    "num_beams": 4,
}


def load_expertqa_for_autoais(data_path):
    ds = read_jsonl(data_path)

    autoais_examples = []
    for ex in ds:
        for system, res in ex["answers"].items():
            for c in res["claims"]:
                if "evidence" not in c or len(c["evidence"]) == 0:
                    continue

                label = "1" if c["support"] == "Complete" else "0"
                claim_str = cleanup_claim_and_evidence(c["claim_string"])
                evidence_str = [cleanup_claim_and_evidence(e) for e in c["evidence"]]
                evidence_str = " ".join(evidence_str)

                input_str = format_example_for_autoais(evidence=evidence_str, claim=claim_str)
                autoais_examples.append({
                    "inputs": input_str,
                    "targets": label
                })

    return autoais_examples


def preprocess_function(sample,
                        tokenizer,
                        padding: str = "max_length",
                        max_seq_len: int = 512,
                        input_column: str = "inputs",
                        target_column: str = "targets"
                        ):
    # created prompted input
    inputs = sample[input_column]
    targets = sample[target_column]

    # tokenize inputs
    model_inputs = tokenizer(inputs, max_length=max_seq_len, padding=padding, truncation=True)

    # Tokenize targets with the `text_target` keyword argument
    labels = tokenizer(text_target=targets, max_length=max_seq_len, padding=padding, truncation=True)

    # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
    # padding in the loss.
    if padding == "max_length":
        labels["input_ids"] = [
            [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
        ]

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs


def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [label.strip() for label in labels]

    # rougeLSum expects newline after each sentence
    preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
    labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]

    return preds, labels


def training_function(args):
    # set seed
    set_seed(args.seed)

    # load dataset from disk and tokenizer
    train_dataset = load_expertqa_for_autoais(args.train_data)
    eval_dataset = load_expertqa_for_autoais(args.val_data)
    train_dataset = Dataset.from_list(train_dataset)
    eval_dataset = Dataset.from_list(eval_dataset)

    # load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(args.model_id)

    preproc_fn = partial(preprocess_function, tokenizer=tokenizer)
    train_dataset = train_dataset.map(preproc_fn, batched=True)
    eval_dataset = eval_dataset.map(preproc_fn, batched=True)


    # load model from the hub
    model = AutoModelForSeq2SeqLM.from_pretrained(
        args.model_id,
        use_cache=False if args.gradient_checkpointing else True,  # this is needed for gradient checkpointing
    )

    # we want to ignore tokenizer pad token in the loss
    label_pad_token_id = -100
    # Data collator
    data_collator = DataCollatorForSeq2Seq(
        tokenizer, model=model, label_pad_token_id=label_pad_token_id, pad_to_multiple_of=8
    )

    # Define compute metrics function
    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        # Replace -100 in the labels as we can't decode them.
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        # Some simple post-processing
        decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

        result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
        result = {k: round(v * 100, 4) for k, v in result.items()}
        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
        result["gen_len"] = np.mean(prediction_lens)
        return result

    # Define training args
    # output_dir = args.repository_id if args.repository_id else args.model_id.split("/")[-1]
    model_name = args.model_id.split("/")[-1]
    if args.output_dir is None:
        output_dir = "out"
    else:
        output_dir = args.output_dir

    model_output_dir = os.path.join(output_dir, model_name)
    log_output_dir = os.path.join(output_dir, "logs")

    training_args = Seq2SeqTrainingArguments(
        output_dir=model_output_dir,
        per_device_train_batch_size=args.per_device_train_batch_size,
        per_device_eval_batch_size=args.per_device_eval_batch_size,
        predict_with_generate=True,
        generation_max_length=args.generation_max_length,
        generation_num_beams=args.generation_num_beams,
        fp16=False,  # T5 overflows with fp16
        bf16=args.bf16,  # Use BF16 if available
        bf16_full_eval=args.bf16,
        learning_rate=args.lr,
        num_train_epochs=args.epochs,
        deepspeed=args.deepspeed,
        gradient_checkpointing=args.gradient_checkpointing,
        # logging & evaluation strategies
        logging_dir=log_output_dir,
        logging_strategy="steps",
        logging_steps=50,
        report_to="wandb",
        evaluation_strategy="steps",
        save_strategy="steps",
        eval_steps=500,
        save_steps=500,
        save_total_limit=1,
        load_best_model_at_end=True,
    )

    # Create Trainer instance
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
    )

    # Start training
    trainer.train()

    # Save our tokenizer and create model card
    tokenizer.save_pretrained(output_dir)
    trainer.create_model_card()
    # Push the results to the hub
    if args.repository_id:
        trainer.push_to_hub()

def parse_arge():
    """Parse the arguments."""
    parser = argparse.ArgumentParser()
    # add model id and dataset path argument
    parser.add_argument("--model_id", type=str, default="google/flan-t5-xl", help="Model id to use for training.")
    parser.add_argument("--train_data", type=str, default="data/expertqa/domain_train.jsonl", help="Path to expertqa training split.")
    parser.add_argument("--val_data", type=str, default="data/expertqa/domain_val.jsonl", help="Path to expertqa validation split.")

    # add training hyperparameters for epochs, batch size, learning rate, and seed
    parser.add_argument("--epochs", type=int, default=3, help="Number of epochs to train for.")
    parser.add_argument("--per_device_train_batch_size", type=int, default=8, help="Batch size to use for training.")
    parser.add_argument("--per_device_eval_batch_size", type=int, default=8, help="Batch size to use for testing.")
    parser.add_argument("--generation_max_length", type=int, default=10, help="Maximum length to use for generation")
    parser.add_argument("--generation_num_beams", type=int, default=3, help="Number of beams to use for generation.")
    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate to use for training.")
    parser.add_argument("--seed", type=int, default=42, help="Seed to use for training.")
    parser.add_argument("--deepspeed", type=str, default=None, help="Path to deepspeed config file.")
    parser.add_argument("--gradient_checkpointing", type=bool, default=True, help="Path to deepspeed config file.")
    parser.add_argument("--output_dir", type=str, default=None, help="Path to model + log output dir")
    parser.add_argument(
        "--bf16",
        action="store_true",
        help="Whether to use bf16.",
    )
    parser.add_argument(
        "--hf_token",
        type=str,
        default=HfFolder.get_token(),
        help="Token to use for uploading models to Hugging Face Hub.",
    )
    args = parser.parse_known_args()
    return args

def main():
    args, _ = parse_arge()
    training_function(args)


if __name__ == "__main__":
    main()
modeling/auto_attribution/finetune_autoais.py
chaitanyamalaviya/ExpertQA
# copied from https://github.com/shmsw25/FActScore/blob/main/factscore/atomic_facts.py

import json
import numpy as np
import re
import functools
import string
import spacy
import sys
import nltk
import openai
from rank_bm25 import BM25Okapi
import os
import time
from nltk.tokenize import sent_tokenize

from load_open_ai_model import OpenAIModel
from dotenv import load_dotenv

load_dotenv()
nltk.download("punkt")

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--input_file", help="input filepath", type=str)
parser.add_argument("--output_file", help="output filepath", type=str)
args = parser.parse_args()


class AtomicFactGenerator(object):
    def __init__(self, demon_dir, gpt3_cache_file=None):
        self.nlp = spacy.load("en_core_web_sm")
        self.is_bio = True
        self.demon_path = os.path.join(demon_dir, "demons.json" if self.is_bio else "demons_complex.json")

        self.openai_lm = OpenAIModel("InstructGPT", cache_file=gpt3_cache_file)

        # get the demos
        with open(self.demon_path, 'r') as f:
            self.demons = json.load(f)

        tokenized_corpus = [doc.split(" ") for doc in self.demons.keys()]
        self.bm25 = BM25Okapi(tokenized_corpus)

    def save_cache(self):
        self.openai_lm.save_cache()

    def run(self, generation, cost_estimate=None):
        """Convert the generation into a set of atomic facts. Return a total words cost if cost_estimate != None."""
        assert isinstance(generation, str), "generation must be a string"
        paragraphs = [para.strip() for para in generation.split("\n") if len(para.strip()) > 0]
        return self.get_atomic_facts_from_paragraph(paragraphs, cost_estimate=cost_estimate)

    def get_atomic_facts_from_paragraph(self, paragraphs, cost_estimate=None):
        sentences = []
        para_breaks = []
        for para_idx, paragraph in enumerate(paragraphs):
            if para_idx > 0 :
                para_breaks.append(len(sentences))

            initials = detect_initials(paragraph)

            curr_sentences = sent_tokenize(paragraph)
            curr_sentences_2 = sent_tokenize(paragraph)

            curr_sentences = fix_sentence_splitter(curr_sentences, initials)
            curr_sentences_2 = fix_sentence_splitter(curr_sentences_2, initials)

            # checking this, just to ensure the crediability of the sentence splitter fixing algorithm
            assert curr_sentences == curr_sentences_2, (paragraph, curr_sentences, curr_sentences_2)

            sentences += curr_sentences

        atoms_or_estimate = self.get_init_atomic_facts_from_sentence([sent for i, sent in enumerate(sentences) if not (not self.is_bio and ( \
                            (i==0 and (sent.startswith("Sure") or sent.startswith("Here are"))) or \
                            (i==len(sentences)-1 and (sent.startswith("Please") or sent.startswith("I hope") or sent.startswith("Here are")))))], cost_estimate=cost_estimate)

        if cost_estimate:
            return atoms_or_estimate
        else:
            atoms = atoms_or_estimate

        atomic_facts_pairs = []
        for i, sent in enumerate(sentences):
            if not self.is_bio and ( \
                (i==0 and (sent.startswith("Sure") or sent.startswith("Here are"))) or \
                (i==len(sentences)-1 and (sent.startswith("Please") or sent.startswith("I hope") or sent.startswith("Here are")))):
                atomic_facts_pairs.append((sent, []))
            elif self.is_bio and sent.startswith("This sentence does not contain any facts"):
                atomic_facts_pairs.append((sent, []))
            elif sent.startswith("Sure") or sent.startswith("Please") or (i==0 and sent.startswith("Here are")):
                atomic_facts_pairs.append((sent, []))
            else:
                atomic_facts_pairs.append((sent, atoms[sent]))

        # postprocess_atomic_facts will fix minor issues from InstructGPT
        # it is supposed to handle sentence splitter issue too, but since here
        # we fixed sentence splitter issue already,
        # the new para_breaks should be identical to the original para_breaks
        if self.is_bio:
            atomic_facts_pairs, para_breaks = postprocess_atomic_facts(atomic_facts_pairs, list(para_breaks), self.nlp)

        return atomic_facts_pairs, para_breaks


    def get_init_atomic_facts_from_sentence(self, sentences, cost_estimate=None):
        """Get the initial atomic facts from the sentences. Return a total words cost if cost_estimate != None."""

        is_bio = self.is_bio
        demons = self.demons

        k = 1 if is_bio else 0
        n = 7 if is_bio else 8

        prompts = []
        prompt_to_sent = {}
        atoms = {}
        for sentence in sentences:
            if sentence in atoms:
                continue
            top_machings = best_demos(sentence, self.bm25, list(demons.keys()), k)
            prompt = ""

            for i in range(n):
                prompt = prompt + "Please breakdown the following sentence into independent facts: {}\n".format(list(demons.keys())[i])
                for fact in demons[list(demons.keys())[i]]:
                    prompt = prompt + "- {}\n".format(fact)
                prompt = prompt + "\n"

            for match in top_machings:
                prompt = prompt + "Please breakdown the following sentence into independent facts: {}\n".format(match)
                for fact in demons[match]:
                    prompt = prompt + "- {}\n".format(fact)
                prompt = prompt + "\n"
            prompt = prompt + "Please breakdown the following sentence into independent facts: {}\n".format(sentence)
            prompts.append(prompt)
            prompt_to_sent[prompt] = sentence

        if cost_estimate:
            total_words_estimate = 0
            for prompt in prompts:
                if cost_estimate == "consider_cache" and (prompt.strip() + "_0") in self.openai_lm.cache_dict:
                    continue
                total_words_estimate += len(prompt.split())
            return total_words_estimate
        else:
            for prompt in prompts:
                output, _ = self.openai_lm.generate(prompt)
                atoms[prompt_to_sent[prompt]] = text_to_sentences(output)

            for key, value in demons.items():
                if key not in atoms:
                    atoms[key] = value

            return atoms


def best_demos(query, bm25, demons_sents, k):
    tokenized_query = query.split(" ")
    top_machings = bm25.get_top_n(tokenized_query, demons_sents, k)
    return top_machings


# transform InstructGPT output into sentences
def text_to_sentences(text):
    sentences = text.split("- ")[1:]
    sentences = [sent.strip()[:-1] if sent.strip()[-1] == '\n' else sent.strip() for sent in sentences]
    if len(sentences) > 0: 
        if sentences[-1][-1] != '.':
            sentences[-1] = sentences[-1] + '.' 
    else:
        sentences = []
    return sentences


def normalize_answer(s):
    """Lower text and remove punctuation, articles and extra whitespace."""
    def remove_articles(text):
        regex = re.compile(r'\b(a|an|the)\b', re.UNICODE)
        return re.sub(regex, ' ', text)
    def white_space_fix(text):
        return ' '.join(text.split())
    def remove_punc(text):
        exclude = set(string.punctuation)
        return ''.join(ch for ch in text if ch not in exclude)
    def lower(text):
        return text.lower()
    return white_space_fix(remove_articles(remove_punc(lower(s))))

MONTHS = ["January", "February", "March", "April", "May", "June", "July", "August", "September", "October", "November", "December"]
MONTHS = [m.lower() for m in MONTHS]

def is_num(text):
    try:
        text = int(text)
        return True
    except Exception:
        return False

def is_date(text):
    text = normalize_answer(text)
    for token in text.split(" "):
        if (not is_num(token)) and token not in MONTHS:
            return False
    return True

def extract_numeric_values(text):
    pattern = r'\b\d+\b'  # regular expression pattern for integers
    numeric_values = re.findall(pattern, text)  # find all numeric values in the text
    return set([value for value in numeric_values])  # convert the values to float and return as a list


def detect_entities(text, nlp):
    doc = nlp(text)
    entities = set()

    def _add_to_entities(text):
        if "-" in text:
            for _text in text.split("-"):
                entities.add(_text.strip())
        else:
            entities.add(text)


    for ent in doc.ents:
        # spacy often has errors with other types of entities
        if ent.label_ in ["DATE", "TIME", "PERCENT", "MONEY", "QUANTITY", "ORDINAL", "CARDINAL"]:

            if is_date(ent.text):
                _add_to_entities(ent.text)
            else:
                for token in ent.text.split():
                    if is_date(token):
                        _add_to_entities(token)
        
    for new_ent in extract_numeric_values(text):
        if not np.any([new_ent in ent for ent in entities]):
            entities.add(new_ent)

    return entities

def postprocess_atomic_facts(_atomic_facts, para_breaks, nlp):

    verbs = ["born.", " appointed.", " characterized.", " described.", " known.", " member.", " advocate.", "served.", "elected."]
    permitted_verbs = ["founding member."]

    atomic_facts = []
    new_atomic_facts = []
    new_para_breaks = []

    for i, (sent, facts) in enumerate(_atomic_facts):
        sent = sent.strip()
        if len(sent.split())==1 and i not in para_breaks and i > 0:
            assert i not in para_breaks
            atomic_facts[-1][0] += " " + sent
            atomic_facts[-1][1] += facts
        else:
            if i in para_breaks:
                new_para_breaks.append(len(atomic_facts))
            atomic_facts.append([sent, facts])

    for i, (sent, facts) in enumerate(atomic_facts):
        entities = detect_entities(sent, nlp)
        covered_entities = set()
        # print (entities)
        new_facts = []
        for i, fact in enumerate(facts):
            if any([fact.endswith(verb) for verb in verbs]) and not any([fact.endswith(verb) for verb in permitted_verbs]):
                if any([fact[:-1] in other_fact for j, other_fact in enumerate(facts) if j != i]):
                    continue
            sent_entities = detect_entities(fact, nlp)
            covered_entities |= set([e for e in sent_entities if e in entities])
            new_entities = sent_entities - entities
            if len(new_entities) > 0:
                do_pass = False
                for new_ent in new_entities:
                    pre_ent = None
                    for ent in entities:
                        if ent.startswith(new_ent):
                            pre_ent = ent
                            break
                    if pre_ent is None:
                        do_pass = True
                        break
                    fact = fact.replace(new_ent, pre_ent)
                    covered_entities.add(pre_ent)
                if do_pass:
                    continue
            if fact in new_facts:
                continue
            new_facts.append(fact)
        try:
            assert entities==covered_entities
        except Exception:
            new_facts = facts # there is a bug in spacy entity linker, so just go with the previous facts

        new_atomic_facts.append((sent, new_facts))

    return new_atomic_facts, new_para_breaks

def is_integer(s):
    try:
        s = int(s)
        return True
    except Exception:
        return False

def detect_initials(text):
    pattern = r"[A-Z]\. ?[A-Z]\."
    match = re.findall(pattern, text)
    return [m for m in match]

def fix_sentence_splitter(curr_sentences, initials):
    for initial in initials:
        if not np.any([initial in sent for sent in curr_sentences]):
            alpha1, alpha2 = [t.strip() for t in initial.split(".") if len(t.strip())>0]
            for i, (sent1, sent2) in enumerate(zip(curr_sentences, curr_sentences[1:])):
                if sent1.endswith(alpha1 + ".") and sent2.startswith(alpha2 + "."):
                    # merge sentence i and i+1
                    curr_sentences = curr_sentences[:i] + [curr_sentences[i] + " " + curr_sentences[i+1]] + curr_sentences[i+2:]
                    break
    sentences = []
    combine_with_previous = None
    for sent_idx, sent in enumerate(curr_sentences):
        if len(sent.split())<=1 and sent_idx==0:
            assert not combine_with_previous
            combine_with_previous = True
            sentences.append(sent)
        elif len(sent.split())<=1:
            assert sent_idx > 0
            sentences[-1] += " " + sent
            combined_with_previous = False
        elif sent[0].isalpha() and not sent[0].isupper() and sent_idx > 0:
            assert sent_idx > 0, curr_sentences
            sentences[-1] += " " + sent
            combine_with_previous = False
        elif combine_with_previous:
            assert sent_idx > 0
            sentences[-1] += " " + sent
            combine_with_previous = False
        else:
            assert not combine_with_previous
            sentences.append(sent)
    return sentences


def main():
    print(os.getenv('OPENAI_API_KEY'))
    generator = AtomicFactGenerator("modeling/fact_score/demos", gpt3_cache_file='.cache/factscore')

    ## load dataset
    examples = example_utils.read_examples(args.input_file)
    for example in examples:
        answer = example.answers[list(example.answers.keys())[0]]
        for claim in answer['claims']:
            claim_string = re.sub(r'\[\d+\]', '', claim['claim_string'])
            if not claim_string.strip():
                claim['atomic_claims'] = []
            atomic_facts = generator.run(claim_string)
            claim['atomic_claims'] = atomic_facts[0][0][1]
        example_utils.write_examples(args.output_file, [example], True)
            

if __name__ == "__main__":
    main()

modeling/fact_score/break_down_to_atomic_claims.py
chaitanyamalaviya/ExpertQA
"""
Compute FActScore F1 scores with respect to human factuality judgements.
"""
import os
import sys
from sklearn.metrics import accuracy_score, f1_score, precision_recall_fscore_support
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
from data_utils.jsonl_utils import read_jsonl


def did_abstain(answer):
    if "The context does not" in answer or \
    "I don't know" in answer or \
    "The context provided does not" in answer or \
    "The context provided doesn't" in answer or \
    "The context doesn't" in answer or \
    "not sure if I understand" in answer or \
    "I'm sorry, but I can't" in answer or \
    "The given context does not" in answer or \
    "I couldn't find any relevant information" in answer:
        return True
    else:
        return False


def is_void_claim(claim):
    # Filter out any claims that contain the following substrings
    filter_out_claims = set(["Hello",
                         "Hello, this is Bing.", 
                         "this is Bing",
                         "Thank you for your question"
                         "glad you're",
                         "glad you",
                         "glad to help you",
                         "glad to hear",
                         "Bing",
                         "Hope this helps",
                         "Hope this is helpful",
                         "I can provide more information based on your answer", 
                         "are you interested in?", 
                         "Are you looking for", 
                         "what aspect are you interested", 
                         "I found some information that might help you", 
                         "I found some information", 
                         "happy to help you", 
                         "I hope this information",
                         "I hope this helps",
                         "I hope this answer helps you",
                         "I can provide more specific information",
                         "I can provide more details",
                         "I searched for your question and found some results.", 
                         "I can help you",
                         "Thank you for your question",
                         "That's an interesting question",
                         "not sure if I understand your question",
                         "I hope this answers your question",
                         "I can provide you with more specific resources",
                         "I can try to answer your question",
                         "I can provide more information",
                         "I can search for more",
                         "I can answer your question",
                         "I can give you more",
                         "I can try to summarise",
                         "I searched for your question and found some information",
                         "I can answer your question",
                         "I can provide you",
                         "I can try to find",
                         "I cannot give you a definitive answer",
                         "I might be able to provide",
                         "please let me know"])
    for f_claim in filter_out_claims:
        if f_claim in claim:
            return True
    return False




def human_corr(input_file):
    label_to_score = {"Definitely correct": 1.0, "Probably correct": 0.5, "Unsure": 0.0,
                      "Likely incorrect": -0.5, "Definitely incorrect": -1.0}
    examples = read_jsonl(input_file)
    threshold = 0.5

    preds = []
    golds = []
    ipreds = []
    igolds = []
    per_sys_preds = {}
    per_sys_golds = {}
    for ex in examples:
        for system, answer in ex["answers"].items():
            if did_abstain(answer["answer_string"]):
                continue

            if system not in per_sys_preds:
                per_sys_preds[system] = []
                per_sys_golds[system] = []

            for c in answer["claims"]:

                if c["correctness"] == None or c["fact_score"] == None or is_void_claim(c["claim_string"]):
                    continue
                fact_score = c["fact_score"]
                factuality_label = label_to_score[c["correctness"]]

                preds.append(fact_score)
                golds.append(factuality_label)
                ipreds.append(1 if fact_score >= threshold else 0)
                igolds.append(1 if " correct" in c["correctness"] else 0)
                per_sys_preds[system].append(fact_score)
                per_sys_golds[system].append(factuality_label)

    print(f"Accuracy: {accuracy_score(ipreds, igolds)}")
    print(f"F1 score averaged: {f1_score(igolds, ipreds, average='micro')}")
    print(f"F1 score breakdown (P/R/F1/support): {precision_recall_fscore_support(igolds, ipreds, average=None)}")
    print(f"Total number of claims: {len(golds)}\n")

    for system, sys_golds in per_sys_golds.items():
        sys_preds = [1 if fact_score >= threshold else 0 for fact_score in per_sys_preds[system]]
        sys_golds_binary = [1 if gold >= 0.5 else 0 for gold in sys_golds]
        print(f"{system}")
        print(f"Accuracy: {accuracy_score(sys_golds_binary, sys_preds)}")
        print(f"F1 score averaged: {f1_score(sys_golds_binary, sys_preds, average='micro')}")
        print(f"F1 score breakdown (P/R/F1/support): {precision_recall_fscore_support(sys_golds_binary, sys_preds, average=None)}")
        print(f"\t\tNumber of claims: {len(sys_golds_binary)}")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file",
                        default="data/r2_compiled_out_corrected_revised_atomic_w_evidences_factscores.jsonl",
                        help="example file with per-claim fact scores",
                        type=str)
    args = parser.parse_args()

    human_corr(args.input_file)

modeling/fact_score/compute_factuality_f1.py
chaitanyamalaviya/ExpertQA
# based off of https://github.com/shmsw25/FActScore/blob/main/factscore/factscorer.py

import argparse
import string
import os
import sys
import numpy as np
from tqdm import tqdm

from load_open_ai_model import OpenAIModel
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils


parser = argparse.ArgumentParser()
parser.add_argument("--input_file", help="input filepath", type=str)
parser.add_argument("--output_file", help="output filepath", type=str)
args = parser.parse_args()


def get_score(topic, atomic_claim, atomic_evidence, lm):
        atomic_claim = atomic_claim.strip()
        # if lm:
        definition = "Answer the question about {} based on the given context.\n\n".format(topic)
        definition += atomic_evidence.strip()
        if not definition[-1] in string.punctuation:
            definition += "."
        prompt = "{}\n\nInput: {} True or False?\nOutput:".format(definition.strip(), atomic_claim.strip())

        output = lm.generate(prompt)

        if type(output[1])==np.ndarray:
            # when logits are available
            logits = np.array(output[1])
            assert logits.shape[0] in [32000, 32001]
            true_score = logits[5852]
            false_score = logits[7700]
            is_supported = true_score > false_score
        else:
            # when logits are unavailable
            generated_answer = output[0].lower()
            if "true" in generated_answer or "false" in generated_answer:
                if "true" in generated_answer and "false" not in generated_answer:
                    is_supported = True
                elif "false" in generated_answer and "true" not in generated_answer:
                    is_supported = False
                else:
                    is_supported = generated_answer.index("true") > generated_answer.index("false")
            else:
                is_supported = all([keyword not in generated_answer.lower().translate(str.maketrans("", "", string.punctuation)).split() for keyword in ["not", "cannot", "unknown", "information"]])

        # else:
        #     is_supported = True
        
        return is_supported


if __name__ == '__main__':

    lm = OpenAIModel("ChatGPT", cache_file='.cache/factscore')

    input_data = example_utils.read_examples(args.input_file)

    for example in tqdm(input_data):
        claims = example.answers[list(example.answers.keys())[0]]["claims"]
        for claim in claims:
            score_total = 0
            for i in range(len(claim["atomic_claims"])):
                a_claim = claim["atomic_claims"][i]
                a_evidence = claim["atomic_evidences"][i]
                if "Passage ID 4" in a_evidence:
                    a_evidence = a_evidence.split("Passage ID 4")[0]
                score = get_score(example.metadata.field, a_claim, a_evidence, lm)
                score_total += score
            if claim["atomic_claims"]:
                claim["fact_score"] = float(score_total / len(claim["atomic_claims"]))
            else:
                claim["fact_score"] = None
        example_utils.write_examples(args.output_file, [example], True)

modeling/fact_score/factscore.py
chaitanyamalaviya/ExpertQA
# copied from https://github.com/shmsw25/FActScore/blob/main/factscore/atomic_facts.py

import pickle
import os
import time

class LM(object):

    def __init__(self, cache_file):
        self.cache_file = cache_file
        self.cache_dict = self.load_cache()
        self.model = None
        self.add_n = 0

    def load_model(self):
        # load the model and put it as self.model
        raise NotImplementedError()

    def generate(self, prompt, sample_idx=0, max_sequence_length=2048, max_output_length=128):
        prompt = prompt.strip() # it's important not to end with a whitespace
        cache_key = f"{prompt}_{sample_idx}"

        if cache_key in self.cache_dict:
            return self.cache_dict[cache_key]

        if self.model is None:
            self.load_model()

        if prompt.endswith(" True or False?\nAnswer:"):
            generated = self._generate(prompt, max_sequence_length=max_sequence_length, max_output_length=1)
        else:
            generated = self._generate(prompt, max_sequence_length=max_sequence_length, max_output_length=max_output_length)

        self.cache_dict[cache_key] = generated
        self.add_n += 1
        return generated

    def save_cache(self):
        if self.add_n == 0:
            return

        # load the latest cache first, since if there were other processes running in parallel, cache might have been updated
        for k, v in self.load_cache().items():
            self.cache_dict[k] = v

        with open(self.cache_file, "wb") as f:
            pickle.dump(self.cache_dict, f)

    def load_cache(self, allow_retry=True):
        if os.path.exists(self.cache_file):
            while True:
                try:
                    with open(self.cache_file, "rb") as f:
                        cache = pickle.load(f)
                    break
                except Exception:
                    if not allow_retry:
                        assert False
                    print ("Pickle Error: Retry in 5sec...")
                    time.sleep(5)        
        else:
            cache = {}
        return cache
modeling/fact_score/load_lm.py
chaitanyamalaviya/ExpertQA
# copied from https://github.com/shmsw25/FActScore/blob/main/factscore/atomic_facts.py

from load_lm import LM
import openai
import sys
import time
import os
import numpy as np
import logging

class OpenAIModel(LM):

    def __init__(self, model_name, cache_file=None):
        self.model_name = model_name
        self.temp = 0.7
        self.save_interval = 100
        super().__init__(cache_file)

    def load_model(self):
        # load api key
        openai.api_key = os.getenv('OPENAI_API_KEY')
        openai.organization = '' # Fill in organization ID
        self.model = self.model_name

    def _generate(self, prompt, max_sequence_length=2048, max_output_length=128):
        if self.add_n % self.save_interval == 0:
            self.save_cache()
        # return a tuple of string (generated text) and metadata (any format)
        # This should be about generating a response from the prompt, no matter what the application is
        if self.model_name == "ChatGPT":
            # Construct the prompt send to ChatGPT
            message = [{"role": "user", "content": prompt}]
            # Call API
            response = call_ChatGPT(message, temp=self.temp, max_len=max_sequence_length)
            # Get the output from the response
            output = response["choices"][0]["message"]["content"]
            return output, response
        elif self.model_name == "InstructGPT":
            # Call API
            response = call_GPT3(prompt, temp=self.temp)
            # Get the output from the response
            output = response["choices"][0]["text"]
            return output, response
        else:
            raise NotImplementedError()

def call_ChatGPT(message, model_name="gpt-3.5-turbo", max_len=1024, temp=0.7, verbose=False):
    # call GPT-3 API until result is provided and then return it
    response = None
    received = False
    num_rate_errors = 0
    while not received:
        try:
            response = openai.ChatCompletion.create(model=model_name,
                                                    messages=message,
                                                    max_tokens=max_len,
                                                    temperature=temp)
            received = True
        except:
            # print(message)
            num_rate_errors += 1
            error = sys.exc_info()[0]
            if error == openai.error.InvalidRequestError:
                # something is wrong: e.g. prompt too long
                logging.critical(f"InvalidRequestError\nPrompt passed in:\n\n{message}\n\n")
                assert False
            
            logging.error("API error: %s (%d). Waiting %dsec" % (error, num_rate_errors, np.power(2, num_rate_errors)))
            time.sleep(np.power(2, num_rate_errors))
    return response


def call_GPT3(prompt, model_name="text-davinci-003", max_len=512, temp=0.7, num_log_probs=0, echo=False, verbose=False):
    # call GPT-3 API until result is provided and then return it
    response = None
    received = False
    num_rate_errors = 0
    while not received:
        try:
            response = openai.Completion.create(model=model_name,
                                                prompt=prompt,
                                                max_tokens=max_len,
                                                temperature=temp,
                                                logprobs=num_log_probs,
                                                echo=echo)
            received = True
        except:
            error = sys.exc_info()[0]
            num_rate_errors += 1
            if error == openai.error.InvalidRequestError:
                # something is wrong: e.g. prompt too long
                logging.critical(f"InvalidRequestError\nPrompt passed in:\n\n{prompt}\n\n")
                assert False
            logging.error("API error: %s (%d)" % (error, num_rate_errors))
            time.sleep(np.power(2, num_rate_errors))
    return response
modeling/fact_score/load_open_ai_model.py
chaitanyamalaviya/ExpertQA
import numpy as np
import os
import requests
import sys
from tqdm import tqdm
from typing import List

from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain import VectorDBQA
from langchain.chains import RetrievalQA, RetrievalQAWithSourcesChain
from langchain.schema import Document
from langchain.prompts import PromptTemplate

from newspaper import Article
from newspaper.article import ArticleException

import openai

from bs4 import BeautifulSoup # Required to parse HTML
from urllib.parse import unquote # Required to unquote URLs


import argparse

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils


parser = argparse.ArgumentParser()
parser.add_argument("--input_file", help="input filepath", type=str)
parser.add_argument("--output_file", help="output filepath", type=str)
parser.add_argument("--topk", help="the value of k for the topk passages to use for QA", type=int, default=5)
parser.add_argument("--start_idx", type=int, default=0)

args = parser.parse_args()


text_splitter = CharacterTextSplitter(separator=' ', chunk_size=1000, chunk_overlap=200)

#####
# OPENAI
#####
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
if OPENAI_API_KEY is None:
    raise ValueError("Please set the OpenAI API key as environment variable 'OPENAI_API_KEY'.")

openai.organization = ""
openai.api_key = ""
#####
# Google Search
#####

def google_search(
        query: str,
        api_key: str = None,
        cx: str = None
):
    """Get top 10 webpages from Google search.

    Args:
        query: search query
        api_key: custom search engine api key
        cx: programmable search engine id
    Returns:
        top-10 search results in json format
    """
    response = requests.get(f"https://www.google.com/search?q={query}") # Make the request
    soup = BeautifulSoup(response.text, "html.parser") # Parse the HTML
    links = soup.find_all("a") # Find all the links in the HTML

    urls = []
    for l in [link for link in links if link["href"].startswith("/url?q=")]:
        # get the url
        url = l["href"]
        # remove the "/url?q=" part
        url = url.replace("/url?q=", "")
        # remove the part after the "&sa=..."
        url = unquote(url.split("&sa=")[0])
        # special case for google scholar
        if url.startswith("https://scholar.google.com/scholar_url?url=http"):
            url = url.replace("https://scholar.google.com/scholar_url?url=", "").split("&")[0]
        elif 'google.com/' in url: # skip google links
            continue
        if url.endswith('.pdf'): # skip pdf links
            continue
        if '#' in url: # remove anchors (e.g. wikipedia.com/bob#history and wikipedia.com/bob#genetics are the same page)
            url = url.split('#')[0]
        # print the url
        urls.append(url)

    # Use numpy to dedupe the list of urls after removing anchors
    urls = list(np.unique(urls))
    return urls

def scrape_and_parse(url: str):
    """Scrape a webpage and parse it into a Document object"""
    a = Article(url)
    try:
        a.download()
        a.parse()
    except Exception as e:
        return None

    return {
        "url": url,
        "text": a.text,
    }


def scrape_and_filter(urls: list):
    doc_list = []
    for u in urls:
        print(f"Processing: {u}")
        doc = scrape_and_parse(u)
        if doc is None:
            continue
        elif "Access" in doc["text"] and "Denied" in doc["text"]:
            continue
        else:
            doc_list.append(doc)

    return doc_list


def retrieve_topk_passages(query: str,
                           retrieved_documents: List[Document],
                           topk: int):

    texts = text_splitter.split_documents(retrieved_documents)
    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY, max_retries=1000)
    while True:
        try:
            docsearch = Chroma.from_documents(texts, embeddings)
        except Exception as e:
            continue
        break
    doc_retriever = docsearch.as_retriever(search_kwargs={"k": topk})
    topk_relevant_passages = doc_retriever.get_relevant_documents(query)
    topk_relevant_passages_content = ["Passage ID " + str(i+1) + ": " + doc.page_content.replace("\n","") for i, doc in enumerate(topk_relevant_passages)]

    return "\n\n".join(topk_relevant_passages_content)


if __name__ == '__main__':

    input_data = example_utils.read_examples(args.input_file)
    start_idx = args.start_idx

    for example in tqdm(input_data[start_idx: start_idx+500]):
        claims = example.answers[list(example.answers.keys())[0]]["claims"]       
        for claim in claims:
            atomic_evidences = []
            for i, a_claim in enumerate(claim["atomic_claims"]):
                # Google search for each atomic claim 
                search_results = google_search(a_claim)
                # search_urls = [r['link'] for r in search_results]
                if len(search_results) > 0:
                    all_docs = scrape_and_filter(search_results)
                else:
                    atomic_evidences.append('')
                    continue
                all_docs = [Document(page_content=d['text'], metadata={'source': d['url']}) for d in all_docs]
                all_docs_content_lens = [len(doc.page_content.strip()) for doc in all_docs]
                if not all_docs or not sum(all_docs_content_lens):
                    atomic_evidences.append('')
                    continue
                relevant_passages = retrieve_topk_passages(a_claim, all_docs, args.topk)
                atomic_evidences.append(relevant_passages)
            claim["atomic_evidences"] = atomic_evidences

        example_utils.write_examples(args.output_file, [example], True)

modeling/fact_score/retrieve_evidence_for_claims.py
chaitanyamalaviya/ExpertQA
'''Script to convert examples for LFQA'''

from absl import app
from absl import flags
import os
import re
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils, jsonl_utils

FLAGS = flags.FLAGS

flags.DEFINE_string('input_prefix', 'data/expertqa/rand_', 'Input filepath.')
flags.DEFINE_string('output_prefix', 'data/expertqa/rand_lfqa_', 'Output filepath.')


def main(unused_argv):
    train_data = example_utils.read_examples(FLAGS.input_prefix + "train.jsonl")
    val_data = example_utils.read_examples(FLAGS.input_prefix + "val.jsonl")
    test_data = example_utils.read_examples(FLAGS.input_prefix + "test.jsonl")
    train_wdata = []
    for i, ex in enumerate(train_data):
        for answer_model_name, answer in ex.answers.items():
            answer_string = re.sub(r'\[\d+\]', '', answer["revised_answer_string"]).replace("\n"," ")
            if not answer_string.strip() or answer_string == None:
                continue
            train_wdata.append({"example_id": len(train_wdata), "context": "", "question": ex.question, "answer": answer_string})
    
    val_wdata = []
    for i, ex in enumerate(val_data):
        for answer_model_name, answer in ex.answers.items():
            answer_string = re.sub(r'\[\d+\]', '', answer["revised_answer_string"]).replace("\n"," ")
            if not answer_string.strip() or answer_string == None:
                continue
            val_wdata.append({"example_id": len(val_wdata), "context": "", "question": ex.question, "answer": answer_string})

    test_wdata = []
    for i, ex in enumerate(test_data):
        for answer_model_name, answer in ex.answers.items():
            # if answer["usefulness"] == "Not useful at all":
            #     continue
            answer_string = re.sub(r'\[\d+\]', '', answer["revised_answer_string"]).replace("\n"," ")
            if not answer_string.strip() or answer_string == None:
                continue
            test_wdata.append({"example_id": len(test_wdata), "context": "", "question": ex.question, "answer": answer_string})
    
    jsonl_utils.write_jsonl(FLAGS.output_prefix + "train.json", train_wdata)
    jsonl_utils.write_jsonl(FLAGS.output_prefix + "val.json", val_wdata)
    jsonl_utils.write_jsonl(FLAGS.output_prefix + "test.json", test_wdata)


if __name__ == "__main__":
    app.run(main)
modeling/lfqa/convert_for_lfqa.py
chaitanyamalaviya/ExpertQA
#!/usr/bin/env python
# coding=utf-8
# Copyright 2021 The HuggingFace Team All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Fine-tuning the library's seq2seq models for question answering using the 🤗 Seq2SeqTrainer.
"""
# You can also adapt this script on your own question answering task. Pointers for this are left as comments.

import logging
import os
import sys
from dataclasses import dataclass, field
from typing import List, Optional, Tuple

import datasets
import evaluate
import numpy as np
from datasets import load_dataset
import wandb
import json
import math
import sys
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset


import transformers
from transformers import (
    AutoConfig,
    LlamaConfig,
    AutoModelForSeq2SeqLM,
    LlamaForCausalLM,
    AutoTokenizer,
    LlamaTokenizer,
    DataCollatorForSeq2Seq,
    HfArgumentParser,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer,
    set_seed,
)
from trl import SFTTrainer
from transformers.trainer_utils import EvalLoopOutput, EvalPrediction, PredictionOutput, get_last_checkpoint, speed_metrics
from transformers.utils import check_min_version, send_example_telemetry
from transformers.utils.versions import require_version

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
# check_min_version("4.30.0.dev0")

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

logger = logging.getLogger(__name__)


def wandb_hp_space(trial):
    return {
        "method": "random",
        "metric": {"name": "best_f1", "goal": "maximize"},
        "parameters": {
            "learning_rate": {'values': [3e-5, 1e-4]},
            "per_device_train_batch_size": {"values":  [1, 2, 8, 16]},
        },
    }

# sweep_configuration = {
#     'method': 'random',
#     'name': 'sweep',
#     'metric': {
#         'goal': 'maximize', 
#         'name': 'val_acc'
#         },
#     'parameters': {
#         'batch_size': {'values': [1, 2, 8, 16]},
#         'lr': {'values': [3e-5, 1e-4]}
#      }
# }

global val_filename

class QuestionAnsweringSeq2SeqTrainer(Seq2SeqTrainer):
    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 evaluate(self, eval_dataset=None, eval_examples=None, ignore_keys=None, metric_key_prefix: str = "eval"):
    def evaluate(
        self,
        eval_dataset: Optional[Dataset] = None,
        eval_examples=None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
        **gen_kwargs,
    ) -> Dict[str, float]:
        gen_kwargs = gen_kwargs.copy()
        gen_kwargs["max_length"] = (
            gen_kwargs["max_length"] if gen_kwargs.get("max_length") is not None else self.args.generation_max_length
        )
        gen_kwargs["num_beams"] = (
            gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.args.generation_num_beams
        )
        self._gen_kwargs = gen_kwargs

        eval_dataset = self.eval_dataset if eval_dataset is None else eval_dataset
        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        eval_examples = self.eval_examples if eval_examples is None else eval_examples

        # Temporarily disable metric computation, we will do it in the loop here.
        compute_metrics = self.compute_metrics
        self.compute_metrics = None
        start_time = time.time()
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
            output = eval_loop(
                eval_dataloader,
                description="Evaluation",
                # No point gathering the predictions if there are no metrics, otherwise we defer to
                # self.args.prediction_loss_only
                prediction_loss_only=True if compute_metrics is None else None,
                # prediction_loss_only=None,
                ignore_keys=ignore_keys,
                metric_key_prefix=metric_key_prefix,
            )
        finally:
            self.compute_metrics = compute_metrics
        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )
        
        preds = np.where(output.predictions != -100, output.predictions, self.tokenizer.pad_token_id)
        text_outputs = self.tokenizer.batch_decode(preds, skip_special_tokens=True, clean_up_tokenization_spaces=True)
        # text_outputs = self.tokenizer.batch_decode(output.predictions, skip_special_tokens=True)
        val_data = []
        with open(val_filename) as fp:
            val_data = [json.loads(line) for line in fp.readlines()]
        for i, out in enumerate(text_outputs):
            val_data[i]["output"] = [out]
            val_data[i]["context_provided"] = self.tokenizer.decode(eval_dataset[i]['input_ids'])
            val_data[i]["context_raw"] = output.predictions.tolist()

        with open(f"{self.args.output_dir}/test_eval_outputs.json", "w") as fp:
            for line in val_data:
                json.dump(line, fp)
                fp.write("\n")

        if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
            # Only the main node write the results by default
            eval_preds = self.post_process_function(eval_examples, eval_dataset, output)
            metrics = self.compute_metrics(eval_preds)

            # Prefix all keys with metric_key_prefix + '_'
            for key in list(metrics.keys()):
                if not key.startswith(f"{metric_key_prefix}_"):
                    metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

            metrics.update(output.metrics)
        else:
            metrics = output.metrics

        if self.args.should_log:
            # Only the main node log the results by default
            self.log(metrics)

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

    def predict(
        self, predict_dataset, predict_examples, ignore_keys=None, metric_key_prefix: str = "test", **gen_kwargs
    ):
        self._gen_kwargs = gen_kwargs.copy()
        gen_kwargs["max_length"] = (
            gen_kwargs["max_length"] if gen_kwargs.get("max_length") is not None else self.args.generation_max_length
        )
        gen_kwargs["num_beams"] = (
            gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.args.generation_num_beams
        )
        self._gen_kwargs = gen_kwargs

        predict_dataloader = self.get_test_dataloader(predict_dataset)

        # Temporarily disable metric computation, we will do it in the loop here.
        compute_metrics = self.compute_metrics
        self.compute_metrics = None
        start_time = time.time()
        eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
        try:
            output = eval_loop(
                predict_dataloader,
                description="Prediction",
                # No point gathering the predictions if there are no metrics, otherwise we defer to
                # self.args.prediction_loss_only
                prediction_loss_only=True if compute_metrics is None else None,
                ignore_keys=ignore_keys,
                metric_key_prefix=metric_key_prefix,
            )
        finally:
            self.compute_metrics = compute_metrics

        total_batch_size = self.args.eval_batch_size * self.args.world_size
        if f"{metric_key_prefix}_jit_compilation_time" in output.metrics:
            start_time += output.metrics[f"{metric_key_prefix}_jit_compilation_time"]
        output.metrics.update(
            speed_metrics(
                metric_key_prefix,
                start_time,
                num_samples=output.num_samples,
                num_steps=math.ceil(output.num_samples / total_batch_size),
            )
        )

        text_outputs = self.tokenizer.batch_decode(output.predictions, skip_special_tokens=True)
        val_data = []
        with open(val_filename) as fp:
            val_data = [json.loads(line) for line in fp.readlines()]
        for i, out in enumerate(text_outputs):
            val_data[i]["output"] = [out]
            val_data[i]["context_provided"] = self.tokenizer.decode(predict_dataset[i]['input_ids'])

        with open(f"{self.args.output_dir}/test_eval_outputs.json", "w") as fp:
            for line in val_data:
                json.dump(line, fp)
                fp.write("\n")

        if self.post_process_function is None or self.compute_metrics is None:
            return output

        predictions = self.post_process_function(predict_examples, predict_dataset, output, "predict")
        metrics = self.compute_metrics(predictions)

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)
        metrics.update(output.metrics)
        return PredictionOutput(predictions=predictions.predictions, label_ids=predictions.label_ids, metrics=metrics)


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    cache_dir: Optional[str] = field(
        default="/nlp/data/huggingface_cache/",
        metadata={"help": "Path to directory to store the pretrained models downloaded from huggingface.co"},
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
    )
    model_revision: str = field(
        default="main",
        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
    )
    use_auth_token: bool = field(
        default=False,
        metadata={
            "help": (
                "Will use the token generated when running `huggingface-cli login` (necessary to use this script "
                "with private models)."
            )
        },
    )


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    dataset_name: Optional[str] = field(
        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
    )
    dataset_config_name: Optional[str] = field(
        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
    )
    context_column: Optional[str] = field(
        default="context",
        metadata={"help": "The name of the column in the datasets containing the contexts (for question answering)."},
    )
    question_column: Optional[str] = field(
        default="question",
        metadata={"help": "The name of the column in the datasets containing the questions (for question answering)."},
    )
    answer_column: Optional[str] = field(
        default="answers",
        metadata={"help": "The name of the column in the datasets containing the answers (for question answering)."},
    )
    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
    validation_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
    )
    test_file: Optional[str] = field(
        default=None,
        metadata={"help": "An optional input test data file to evaluate the perplexity on (a text file)."},
    )
    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": "The number of processes to use for the preprocessing."},
    )
    max_seq_length: int = field(
        default=512,
        metadata={
            "help": (
                "The maximum total input sequence length after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded."
            )
        },
    )
    max_answer_length: int = field(
        default=512,
        metadata={
            "help": (
                "The maximum length of an answer that can be generated. This is needed because the start "
                "and end predictions are not conditioned on one another."
            )
        },
    )
    val_max_answer_length: Optional[int] = field(
        default=512,
        metadata={
            "help": (
                "The maximum total sequence length for validation target text after tokenization. Sequences longer "
                "than this will be truncated, sequences shorter will be padded. Will default to `max_answer_length`."
                "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
                "during ``evaluate`` and ``predict``."
            )
        },
    )
    pad_to_max_length: bool = field(
        default=True,
        metadata={
            "help": (
                "Whether to pad all samples 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 purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    max_predict_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of prediction examples to this "
                "value if set."
            )
        },
    )
    version_2_with_negative: bool = field(
        default=False, metadata={"help": "If true, some of the examples do not have an answer."}
    )
    null_score_diff_threshold: float = field(
        default=0.0,
        metadata={
            "help": (
                "The threshold used to select the null answer: if the best answer has a score that is less than "
                "the score of the null answer minus this threshold, the null answer is selected for this example. "
                "Only useful when `version_2_with_negative=True`."
            )
        },
    )
    doc_stride: int = field(
        default=128,
        metadata={"help": "When splitting up a long document into chunks, how much stride to take between chunks."},
    )
    n_best_size: int = field(
        default=20,
        metadata={"help": "The total number of n-best predictions to generate when looking for an answer."},
    )
    num_beams: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
                "which is used during ``evaluate`` and ``predict``."
            )
        },
    )
    ignore_pad_token_for_loss: bool = field(
        default=True,
        metadata={
            "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
        },
    )

    def __post_init__(self):
        if (
            self.dataset_name is None
            and self.train_file is None
            and self.validation_file is None
            and self.test_file is None
        ):
            raise ValueError("Need either a dataset name or a training/validation file/test_file.")
        else:
            if self.train_file is not None:
                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 not None:
                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 not None:
                extension = self.test_file.split(".")[-1]
                assert extension in ["csv", "json"], "`test_file` should be a csv or a json file."
        if self.val_max_answer_length is None:
            self.val_max_answer_length = self.max_answer_length


question_answering_column_name_mapping = {
    "squad_v2": ("question", "context", "answer"),
    "nq": ("question", "context", "answer"),
    "drop": ("question", "context", "answer")
}


def main():
    # See all possible arguments in src/transformers/training_args.py
    # or by passing the --help flag to this script.
    # We now keep distinct sets of args, for a cleaner separation of concerns.

    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
    # information sent is the one passed as arguments along with your Python/PyTorch versions.
    send_example_telemetry("run_seq2seq_qa", model_args, data_args)

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

    if training_args.should_log:
        # The default of training_args.log_level is passive, so we set log level at info here to have that default.
        transformers.utils.logging.set_verbosity_info()

    log_level = training_args.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()

    # Log on each process the small summary:
    logger.warning(
        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
    )
    logger.info(f"Training/evaluation parameters {training_args}")

    # Detecting last checkpoint.
    last_checkpoint = None
    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
        last_checkpoint = get_last_checkpoint(training_args.output_dir)
        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
            raise ValueError(
                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
                "Use --overwrite_output_dir to overcome."
            )
        elif last_checkpoint is not None and training_args.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 before initializing model.
    set_seed(training_args.seed)
    global val_filename
    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
    # (the dataset will be downloaded automatically from the datasets Hub).
    #
    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
    # 'text' is found. You can easily tweak this behavior (see below).
    #
    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
    # download the dataset.
    if data_args.dataset_name is not None:
        # Downloading and loading a dataset from the hub.
        raw_datasets = load_dataset(
            data_args.dataset_name,
            data_args.dataset_config_name,
            cache_dir=model_args.cache_dir,
            use_auth_token=True if model_args.use_auth_token else None,
        )
    else:
        data_files = {}
        if data_args.train_file is not None:
            data_files["train"] = data_args.train_file
            extension = data_args.train_file.split(".")[-1]
        if data_args.validation_file is not None:
            data_files["validation"] = data_args.validation_file
            val_filename = data_args.validation_file
            extension = data_args.validation_file.split(".")[-1]
        if data_args.test_file is not None:
            data_files["test"] = data_args.test_file
            extension = data_args.test_file.split(".")[-1]
        # raw_datasets = load_dataset(
        #     extension,
        #     data_files=data_files,
        #     field="data",
        #     cache_dir=model_args.cache_dir,
        #     use_auth_token=True if model_args.use_auth_token else None,
        # )
        raw_datasets = load_dataset(extension, data_files=data_files, cache_dir="/mnt/nlpgridio3/data/cmalaviya/huggingface_data/")
    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
    # https://huggingface.co/docs/datasets/loading_datasets.html.

    # Load pretrained model and tokenizer
    #
    # Distributed training:
    # The .from_pretrained methods guarantee that only one local process can concurrently
    # download model & vocab.

    config = AutoConfig.from_pretrained(
        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
        cache_dir=model_args.cache_dir,
        use_fast=model_args.use_fast_tokenizer,
        revision=model_args.model_revision,
        use_auth_token=True if model_args.use_auth_token else None,
    )
    model = AutoModelForSeq2SeqLM.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=True if model_args.use_auth_token else None,
    )

    # We resize the embeddings only when necessary to avoid index errors. If you are creating a model from scratch
    # on a small vocab and want a smaller embedding size, remove this test.
    embedding_size = model.get_input_embeddings().weight.shape[0]
    if len(tokenizer) > embedding_size:
        model.resize_token_embeddings(len(tokenizer))

    # if model.config.decoder_start_token_id is None:
    #     raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")

    # Preprocessing the datasets.
    # We need to generate and tokenize inputs and targets.
    if training_args.do_train:
        column_names = raw_datasets["train"].column_names
    elif training_args.do_eval:
        column_names = raw_datasets["validation"].column_names
    elif training_args.do_predict:
        column_names = raw_datasets["test"].column_names
    else:
        logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.")
        return

    # Get the column names for input/target.
    dataset_columns = question_answering_column_name_mapping.get(data_args.dataset_name, None)
    if data_args.question_column is None:
        question_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
    else:
        question_column = data_args.question_column
        if question_column not in column_names:
            raise ValueError(
                f"--question_column' value '{data_args.question_column}' needs to be one of: {', '.join(column_names)}"
            )
    if data_args.context_column is None:
        context_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
    else:
        context_column = data_args.context_column
        if context_column not in column_names:
            raise ValueError(
                f"--context_column' value '{data_args.context_column}' needs to be one of: {', '.join(column_names)}"
            )
    if data_args.answer_column is None:
        answer_column = dataset_columns[2] if dataset_columns is not None else column_names[2]
    else:
        answer_column = data_args.answer_column
        if answer_column not in column_names:
            raise ValueError(
                f"--answer_column' value '{data_args.answer_column}' needs to be one of: {', '.join(column_names)}"
            )

    # Temporarily set max_answer_length for training.
    max_answer_length = data_args.max_answer_length
    padding = "max_length" if data_args.pad_to_max_length else False

    if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
        logger.warning(
            "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
            f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
        )

    if data_args.max_seq_length > tokenizer.model_max_length:
        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 preprocess_squad_batch(
        examples,
        question_column: str,
        context_column: str,
        answer_column: str,
    ) -> Tuple[List[str], List[str]]:
        questions = examples[question_column]
        contexts = examples[context_column]
        answers = examples[answer_column]

        def generate_input(_question, _context):
            # return " ".join(["question:", _question.lstrip(), "context:", _context.lstrip()])
            return _question.lstrip()

        inputs = [generate_input(question, context) for question, context in zip(questions, contexts)]
        # targets = [answer["text"][0] if len(answer["text"]) > 0 else "" for answer in answers]
        targets = answers[:]
        return inputs, targets

    def preprocess_function(examples):
        inputs, targets = preprocess_squad_batch(examples, question_column, context_column, answer_column)

        model_inputs = tokenizer(inputs, max_length=max_seq_length, padding=padding, truncation=True)
        # Tokenize targets with text_target=...
        labels = tokenizer(text_target=targets, max_length=max_answer_length, padding=padding, truncation=True)

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]

        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    # Validation preprocessing
    def preprocess_validation_function(examples):
        inputs, targets = preprocess_squad_batch(examples, question_column, context_column, answer_column)

        model_inputs = tokenizer(
            inputs,
            max_length=max_seq_length,
            padding=padding,
            truncation=True,
            # return_overflowing_tokens=True,
            # return_offsets_mapping=True,
        )
        # Tokenize targets with the `text_target` keyword argument
        labels = tokenizer(text_target=targets, max_length=max_answer_length, padding=padding, truncation=True)

        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
        # padding in the loss.
        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
            labels["input_ids"] = [
                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
            ]

        # Since one example might give us several features if it has a long context, we need a map from a feature to
        # its corresponding example. This key gives us just that.
        # sample_mapping = model_inputs.pop("overflow_to_sample_mapping")

        # For evaluation, we will need to convert our predictions to substrings of the context, so we keep the
        # corresponding example_id and we will store the offset mappings.
        model_inputs["example_id"] = []
        # Augment the overflowing tokens to the labels
        labels_out = []

        for i in range(len(model_inputs["input_ids"])):
            model_inputs["example_id"].append(examples["example_id"][i])
            labels_out.append(labels["input_ids"][i])
            # One example can give several spans, this is the index of the example containing this span of text.
            # sample_index = sample_mapping[i]
            # model_inputs["example_id"].append(examples["example_id"])
            # labels_out.append(labels["input_ids"][sample_index])

        model_inputs["labels"] = labels_out
        return model_inputs


    if training_args.do_train:
        if "train" not in raw_datasets:
            raise ValueError("--do_train requires a train dataset")
        train_dataset = raw_datasets["train"]
        if data_args.max_train_samples is not None:
            # We will select sample from whole data if agument is specified
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))
        # Create train feature from dataset
        with training_args.main_process_first(desc="train dataset map pre-processing"):
            train_dataset = train_dataset.map(
                preprocess_function,
                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 not None:
            # Number of samples might increase during Feature Creation, We select only specified max samples
            max_train_samples = min(len(train_dataset), data_args.max_train_samples)
            train_dataset = train_dataset.select(range(max_train_samples))

    if training_args.do_eval:
        if "validation" not in raw_datasets:
            raise ValueError("--do_eval requires a validation dataset")
        eval_examples = raw_datasets["validation"]
        if data_args.max_eval_samples is not None:
            # We will select sample from whole data
            max_eval_samples = min(len(eval_examples), data_args.max_eval_samples)
            eval_examples = eval_examples.select(range(max_eval_samples))
        # Validation Feature Creation
        with training_args.main_process_first(desc="validation dataset map pre-processing"):
            eval_dataset = eval_examples.map(
                preprocess_validation_function,
                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 not None:
            # During Feature creation dataset samples might increase, we will select required samples again
            max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples)
            eval_dataset = eval_dataset.select(range(max_eval_samples))

    if training_args.do_predict:
        if "test" not in raw_datasets:
            raise ValueError("--do_predict requires a test dataset")
        predict_examples = raw_datasets["test"]
        if data_args.max_predict_samples is not None:
            # We will select sample from whole data
            predict_examples = predict_examples.select(range(data_args.max_predict_samples))
        # Predict Feature Creation
        with training_args.main_process_first(desc="prediction dataset map pre-processing"):
            predict_dataset = predict_examples.map(
                preprocess_validation_function,
                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 not None:
            # During Feature creation dataset samples might increase, we will select required samples again
            max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples)
            predict_dataset = predict_dataset.select(range(max_predict_samples))

    # Data collator
    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
    data_collator = DataCollatorForSeq2Seq(
        tokenizer,
        model=model,
        label_pad_token_id=label_pad_token_id,
        pad_to_multiple_of=8 if training_args.fp16 else None,
    )

    # metric = evaluate.load("squad_v2" if data_args.version_2_with_negative else "squad")
    metric = evaluate.load("rouge")

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

    # Post-processing:
    def post_processing_function(
        examples: datasets.Dataset, features: datasets.Dataset, outputs: EvalLoopOutput, stage="eval"
    ):
        # Decode the predicted tokens.
        preds = outputs.predictions
        if isinstance(preds, tuple):
            preds = preds[0]
        # Replace -100s used for padding as we can't decode them
        preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)

        # Build a map example to its corresponding features.
        example_id_to_index = {k: i for i, k in enumerate(examples["example_id"])}
        feature_per_example = {example_id_to_index[feature["example_id"]]: i for i, feature in enumerate(features)}
        predictions = {}
        # Let's loop over all the examples!
        for example_index, example in enumerate(examples):
            # This is the index of the feature associated to the current example.
            feature_index = feature_per_example[example_index]
            predictions[example["example_id"]] = decoded_preds[feature_index]

        # Format the result to the format the metric expects.
        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": str(k), "prediction_text": v} for k, v in predictions.items()]

        # references = [{"id": ex["example_id"], "answers": ex[answer_column]} for ex in examples]
        references = [{"id": str(ex["example_id"]), "answers": {"text": [ex[answer_column]], "answer_start": [0]}} for ex in examples]
        return EvalPrediction(predictions=formatted_predictions, label_ids=references)

    # Initialize our Trainer
    trainer = QuestionAnsweringSeq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        eval_examples=eval_examples if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
        post_process_function=post_processing_function,
    )

    # Training
    if training_args.do_train:
        checkpoint = None
        if training_args.resume_from_checkpoint is not None:
            checkpoint = training_args.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

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

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

    # Evaluation
    results = {}
    max_length = (
        training_args.generation_max_length
        if training_args.generation_max_length is not None
        else data_args.val_max_answer_length
    )
    num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
    if training_args.do_eval:
        logger.info("*** Evaluate ***")
        metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")

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

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

    # Prediction
    if training_args.do_predict:
        logger.info("*** Predict ***")
        results = trainer.predict(predict_dataset, predict_examples)
        metrics = results.metrics

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

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

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

        trainer.push_to_hub(**kwargs)


def _mp_fn(index):
    # For xla_spawn (TPUs)
    main()


if __name__ == "__main__":
    main()

modeling/lfqa/run_gen_qa.py
chaitanyamalaviya/ExpertQA
# coding=utf-8
# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass, field
from typing import Optional

import os
import torch
from datasets import load_dataset
from peft import LoraConfig
from tqdm import tqdm
from transformers import AutoModelForCausalLM, BitsAndBytesConfig, HfArgumentParser, TrainingArguments, AutoConfig, AutoTokenizer
from accelerate import Accelerator
from peft import PeftModel    

from trl import SFTTrainer
from trl.trainer import ConstantLengthDataset
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils, jsonl_utils

tqdm.pandas()


llama2_prompt = "<s>[INST] <<SYS>>\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n<</SYS>>\n\n[QUESTION] [/INST]"
alpaca_prompt = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\nAnswer the following question completely and accurately.\n\n### Input:\n[QUESTION]\n\n### Response:"
vicuna_prompt = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n\nUSER: [QUESTION]\nASSISTANT:"


# Define and parse arguments.
@dataclass
class ScriptArguments:
    """
    The name of the Casual LM model we wish to fine with SFTTrainer
    """

    model_name: Optional[str] = field(default="llama/Llama-2-7b-hf", metadata={"help": "the model name"})
    dataset_name: Optional[str] = field(
        default="timdettmers/openassistant-guanaco", metadata={"help": "the dataset name"}
    )
    train_file: Optional[str] = field(
        default="data/expertqa/rand_lfqa_train.json", metadata={"help": "train file"}
    )
    validation_file: Optional[str] = field(
        default="data/expertqa/rand_lfqa_val.json", metadata={"help": "validation file"}
    )
    test_file: Optional[str] = field(
        default="data/expertqa/rand_lfqa_test.json", metadata={"help": "test file"}
    )
    dataset_text_field: Optional[str] = field(default="text", metadata={"help": "the text field of the dataset"})
    log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
    learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
    learning_rate_scheduler: Optional[str] = field(default="linear", metadata={"help": "the learning rate scheduler"})
    batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
    seq_length: Optional[int] = field(default=1024, metadata={"help": "Input sequence length"})
    gradient_accumulation_steps: Optional[int] = field(
        default=1, metadata={"help": "the number of gradient accumulation steps"}
    )
    load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 8 bits precision"})
    load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
    use_peft: Optional[bool] = field(default=False, metadata={"help": "Wether to use PEFT or not to train adapters"})
    trust_remote_code: Optional[bool] = field(default=True, metadata={"help": "Enable `trust_remote_code`"})
    output_dir: Optional[str] = field(default="output", metadata={"help": "the output directory"})
    peft_lora_r: Optional[int] = field(default=64, metadata={"help": "the r parameter of the LoRA adapters"})
    peft_lora_alpha: Optional[int] = field(default=16, metadata={"help": "the alpha parameter of the LoRA adapters"})
    logging_steps: Optional[int] = field(default=1, metadata={"help": "the number of logging steps"})
    use_auth_token: Optional[bool] = field(default=True, metadata={"help": "Use HF auth token to access the model"})
    num_train_epochs: Optional[int] = field(default=3, metadata={"help": "the number of training epochs"})
    max_steps: Optional[int] = field(default=-1, metadata={"help": "the number of training steps"})
    do_train: Optional[bool] = field(default=True, metadata={"help": "train the model on train_file"})
    do_eval: Optional[bool] = field(default=True, metadata={"help": "evaluate on validation_file"})

    # parser.add_argument("--eos_token_id", type=int, default=49152)

    # parser.add_argument("--learning_rate", type=float, default=1e-4)
    # parser.add_argument("--lr_scheduler_type", type=str, default="cosine")
    # parser.add_argument("--num_warmup_steps", type=int, default=100)
    # parser.add_argument("--weight_decay", type=float, default=0.05)

    # parser.add_argument("--local_rank", type=int, default=0)
    # parser.add_argument("--no_fp16", action="store_false")
    # parser.add_argument("--bf16", action="store_true", default=False)
    # parser.add_argument("--no_gradient_checkpointing", action="store_false", default=False)
    # parser.add_argument("--seed", type=int, default=0)
    # parser.add_argument("--num_workers", type=int, default=None)
    # parser.add_argument("--output_dir", type=str, default="./checkpoints")
    # parser.add_argument("--log_freq", default=1, type=int)
    # parser.add_argument("--eval_freq", default=1000, type=int)
    # parser.add_argument("--save_freq", default=1000, type=int)


def chars_token_ratio(dataset, tokenizer, nb_examples=400):
    """
    Estimate the average number of characters per token in the dataset.
    """
    total_characters, total_tokens = 0, 0
    for _, example in tqdm(zip(range(nb_examples), iter(dataset)), total=nb_examples):
        text = prepare_sample_text(example)
        total_characters += len(text)
        if tokenizer.is_fast:
            total_tokens += len(tokenizer(text).tokens())
        else:
            total_tokens += len(tokenizer.tokenize(text))

    return total_characters / total_tokens


def prepare_sample_text(example):
    """Prepare the text from a sample of the dataset."""
    # text = f"### Question: {example['question']}\n### Answer: {example['answer']}"
    text = llama2_prompt.replace("[QUESTION]", example["question"]) + f"\n{example['answer']}</s>"
    return text


def create_datasets(tokenizer, script_args):
    data_files = {}
    if script_args.train_file is not None:
        data_files["train"] = script_args.train_file
        extension = script_args.train_file.split(".")[-1]
    if script_args.validation_file is not None:
        data_files["validation"] = script_args.validation_file
        val_filename = script_args.validation_file
        extension = script_args.validation_file.split(".")[-1]
    if script_args.test_file is not None:
        data_files["test"] = script_args.test_file
        extension = script_args.test_file.split(".")[-1]
    raw_datasets = load_dataset(extension, data_files=data_files, cache_dir="/mnt/nlpgridio3/data/cmalaviya/huggingface_data/")

    print(f"Size of the train set: {len(raw_datasets['train'])}. Size of the validation set: {len(raw_datasets['validation'])}")

    chars_per_token = chars_token_ratio(raw_datasets['train'], tokenizer)
    print(f"The character to token ratio of the dataset is: {chars_per_token:.2f}")

    # TODO: Do we really need ConstantLengthDataset?
    
    train_dataset = ConstantLengthDataset(
        tokenizer,
        raw_datasets['train'],
        formatting_func=prepare_sample_text,
        infinite=True,
        seq_length=script_args.seq_length,
        chars_per_token=chars_per_token,
    )
    valid_dataset = ConstantLengthDataset(
        tokenizer,
        raw_datasets['validation'],
        formatting_func=prepare_sample_text,
        infinite=False,
        seq_length=script_args.seq_length,
        chars_per_token=chars_per_token,
    )
    test_dataset = ConstantLengthDataset(
        tokenizer,
        raw_datasets['test'],
        formatting_func=prepare_sample_text,
        infinite=False,
        seq_length=script_args.seq_length,
        chars_per_token=chars_per_token,
    )
    return train_dataset, valid_dataset, test_dataset


def main():

    parser = HfArgumentParser(ScriptArguments)
    script_args = parser.parse_args_into_dataclasses()[0]

    # Step 1: Load the model
    if script_args.load_in_8bit and script_args.load_in_4bit:
        raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
    elif script_args.load_in_8bit or script_args.load_in_4bit:
        quantization_config = BitsAndBytesConfig(
            load_in_8bit=script_args.load_in_8bit, load_in_4bit=script_args.load_in_4bit
        )
        # This means: fit the entire model on the GPU:0
        device_map = {"": 0}
        torch_dtype = torch.bfloat16
    else:
        device_map = None
        quantization_config = None
        torch_dtype = None

    # config = AutoConfig.from_pretrained(script_args.model_name)
    # architecture = config.architectures[0]
    tokenizer = AutoTokenizer.from_pretrained(script_args.model_name, use_auth_token=True)
    print("Setting EOS, BOS, and UNK tokens")
    tokenizer.add_special_tokens(
        {
            "eos_token": "</s>",
            "bos_token": "</s>",
            "unk_token": "</s>",
        }
    )
    print("Loading the model")
    model = AutoModelForCausalLM.from_pretrained(
        script_args.model_name,
        quantization_config=quantization_config,
        device_map=device_map,
        trust_remote_code=script_args.trust_remote_code,
        torch_dtype=torch_dtype,
        use_auth_token=script_args.use_auth_token,
    )

    # Step 2: Load the dataset
    # dataset = load_dataset(script_args.dataset_name, split="train")
    train_dataset, val_dataset, test_dataset = create_datasets(tokenizer, script_args)

    if script_args.do_train:
        # Step 3: Define the training arguments
        training_args = TrainingArguments(
            output_dir=script_args.output_dir,
            per_device_train_batch_size=script_args.batch_size,
            per_device_eval_batch_size=script_args.batch_size,
            gradient_accumulation_steps=script_args.gradient_accumulation_steps,
            learning_rate=script_args.learning_rate,
            logging_steps=script_args.logging_steps,
            num_train_epochs=script_args.num_train_epochs,
            max_steps=script_args.max_steps,
        )
        # training_args = TrainingArguments(
        #     output_dir=args.output_dir,
        #     dataloader_drop_last=True,
        #     evaluation_strategy="steps",
        #     max_steps=args.max_steps,
        #     eval_steps=args.eval_freq,
        #     save_steps=args.save_freq,
        #     logging_steps=args.log_freq,
        #     per_device_train_batch_size=args.batch_size,
        #     per_device_eval_batch_size=args.batch_size,
        #     learning_rate=args.learning_rate,
        #     lr_scheduler_type=args.lr_scheduler_type,
        #     warmup_steps=args.num_warmup_steps,
        #     gradient_accumulation_steps=args.gradient_accumulation_steps,
        #     gradient_checkpointing=not args.no_gradient_checkpointing,
        #     fp16=not args.no_fp16,
        #     bf16=args.bf16,
        #     weight_decay=args.weight_decay,
        #     run_name="llama-7b-finetuned",
        #     report_to="wandb",
        #     ddp_find_unused_parameters=False,
        # )

        # Step 4: Define the LoraConfig
        if script_args.use_peft:
            peft_config = LoraConfig(
                r=script_args.peft_lora_r,
                lora_alpha=script_args.peft_lora_alpha,
                bias="none",
                task_type="CAUSAL_LM",
            )
        else:
            peft_config = None

        # Step 5: Define the Trainer
        trainer = SFTTrainer(
            model=model,
            args=training_args,
            max_seq_length=script_args.seq_length,
            train_dataset=train_dataset,
            eval_dataset=val_dataset,
            # dataset_text_field=script_args.dataset_text_field,
            peft_config=peft_config,
            packing=True
        )
        print("Training...")
        trainer.train()

        # Step 6: Save the model
        print("Saving the model")
        trainer.save_model(script_args.output_dir)
        # trainer.model.save_pretrained(os.path.join(script_args.output_dir, "final_checkpoint/"))

    if script_args.do_eval:
        adapters_name = script_args.output_dir
        model = PeftModel.from_pretrained(model, adapters_name)
        model = model.merge_and_unload()
        model = model.cuda()
        validation_examples = jsonl_utils.read_jsonl(script_args.validation_file)
        for ex in tqdm(validation_examples):
            cur_prompt = llama2_prompt.replace("[QUESTION]", ex["question"])
            input_ids = tokenizer(cur_prompt, return_tensors="pt", truncation=True, max_length=4096).input_ids.cuda()
            outputs = model.generate(input_ids, max_new_tokens=4096, do_sample=True, top_p=0.9, temperature=0.9)
            predicted_ans = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(cur_prompt):]            
            ex["output"] = [predicted_ans]
        
        jsonl_utils.write_jsonl(os.path.join(script_args.output_dir, "test_predictions.jsonl"), validation_examples)

        # predictions = trainer.predict(test_dataset=val_dataset, metric_key_prefix="val")


if __name__ == "__main__":
    main()
modeling/lfqa/run_sft_qa.py
chaitanyamalaviya/ExpertQA
from absl import app
from absl import flags
import json
import os
import sys

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils


def find_urls(citations):
    citations = citations.split()
    urls = []
    MAX_URLS = 100
    for i in range(1, MAX_URLS):
        if f"[{i}]:" in citations:
            idx = citations.index(f"[{i}]:")
            urls.append(f"[{i}] " + citations[idx+1])
        else:
            break
    return urls


def main(unused_argv):
    input_data = example_utils.read_examples("data/r1_data.jsonl")

    with open("data/r1_data_answers_bingchat_balanced_compiled.jsonl") as f:
        bing_responses = [json.loads(line) for line in f.readlines()]
    
    cc = 0
    for ex, resp in zip(input_data, bing_responses):
        # ans_text = ans['item']['messages'][-1]['text']
        # source_urls = [url["seeMoreUrl"] for url in ans['item']['messages'][-1]['sourceAttributions']]
        j = 0
        if "text" not in resp['item']['messages'][-1]['adaptiveCards'][0]['body'][0]:
            j += 1
        if len(resp['item']['messages'][-1]['adaptiveCards'][0]['body']) <= j:
             continue
        vcard = resp['item']['messages'][-1]['adaptiveCards'][0]['body'][j]['text']
        if "\n\n" not in vcard:
            citations = []
        else:
            citations, ans_text = vcard.split("\n\n", 1)
            citations = citations.replace("\n", " ")
            source_urls = find_urls(citations)
        ex.answers["bing_chat"] = example_utils.Answer(answer_string=ans_text,
                                                       attribution=source_urls)
        cc += 1

    example_utils.write_examples("data/r1_data_answers_bingchat_balanced_final.jsonl", input_data)
    print(cc)


if __name__ == "__main__":
    app.run(main)
modeling/response_collection/convert_bingchat_output_to_examples.py
chaitanyamalaviya/ExpertQA
'''Script to get retrieve BingChat responses for questions, that uses https://github.com/acheong08/EdgeGPT. 
You would need to request access to BingChat, then follow directions in the above repo to get `cookies.json`.
'''

import asyncio
import json
from EdgeGPT import Chatbot, ConversationStyle
from absl import app
from absl import flags
from tqdm import tqdm
import os
import sys
import time

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

FLAGS = flags.FLAGS

flags.DEFINE_string('input_file', 'data/r1_data.jsonl', 'Input filepath.')
flags.DEFINE_string('output_file', 'data/r1_data_answers_bingchat_balanced_part2.jsonl', 'Output filepath.')
flags.DEFINE_integer('limit', None, 'Limit to the number of examples processed.')
flags.DEFINE_integer('start_idx', 0, 'Start index.')


async def get_response(question):
    # cookies = json.loads(open("data/cookies.json", encoding="utf-8").read())
    # bot = await Chatbot.create(cookies=cookies)
    bot = await Chatbot.create()
    resp = await bot.ask(prompt=question, conversation_style=ConversationStyle.balanced)
    await bot.close()
    return resp


def find_urls(citations):
    citations = citations.split()
    urls = []
    MAX_URLS = 100
    for i in range(1, MAX_URLS):
        if f"[{i}]:" in citations:
            idx = citations.index(f"[{i}]:")
            urls.append(citations[idx+1])
        else:
            break
    return urls


def main(unused_argv):
    input_data = example_utils.read_examples(FLAGS.input_file)
    sampled_data = input_data[FLAGS.start_idx:]
    fw = open(FLAGS.output_file, "a")
    for i, ex in enumerate(tqdm(sampled_data)):
        if i==FLAGS.limit:
            break
        query = f"I am an expert from the field of {ex.metadata.field}. Please answer my question: " + ex.question
        ans = asyncio.run(get_response(query))
        # if "messages" not in ans['item']:
        #     continue
        # if "text" not in ans['item']['messages'][-1]['adaptiveCards'][0]['body'][0]:
        #     continue
        # ans_text = ans['item']['messages'][-1]['text']
        # source_urls = [url["seeMoreUrl"] for url in ans['item']['messages'][-1]['sourceAttributions']]
       
        # vcard = ans['item']['messages'][-1]['adaptiveCards'][0]['body'][0]['text']
        # citations, ans_text = vcard.split("\n\n", 1)
        # citations = citations.replace("\n", " ")
        # source_urls = find_urls(citations)
        # ex.answers["bing_chat"] = example_utils.Answer(answer_string=ans_text,
        #                                                attribution=source_urls)
        json.dump(ans, fw)
        fw.write("\n")
        time.sleep(3)

    # example_utils.write_examples(FLAGS.output_file, sampled_data)


if __name__ == "__main__":
    app.run(main)
modeling/response_collection/fetch_bingchat_responses.py
chaitanyamalaviya/ExpertQA
"""Given an answer, split it into claims using GPT3"""

import openai
from tqdm import tqdm
from absl import app
from absl import flags
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)
import dataclasses
import json
import os
import sys

# import config

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

openai.organization = ""
openai.api_key = ""

FLAGS = flags.FLAGS

flags.DEFINE_string('input_file', 'data/r1_data.jsonl', 'Input filepath.')
flags.DEFINE_string('output_file', 'data/r1_data_answers_gpt4.jsonl', 'Output filepath.')
flags.DEFINE_string('prompt_config', 'prompts/qa_prompt.txt', 'Prompt path.')
flags.DEFINE_string('model_name', 'gpt4', 'Prompt path.')
flags.DEFINE_integer('limit', None, 'Limit to the number of examples processed')


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def chat_completion_with_backoff(**kwargs):
    return openai.ChatCompletion.create(**kwargs)


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def completion_with_backoff(**kwargs):
    return openai.Completion.create(**kwargs)


def main(unused_argv):

    with open(FLAGS.prompt_config) as f:
        prompt = f.read()

    input_data = example_utils.read_examples(FLAGS.input_file)

    f = open(FLAGS.output_file, "w")

    for i, ex in enumerate(tqdm(input_data)):
        if i==FLAGS.limit:
            break
        cur_prompt = prompt[:]
        cur_prompt = cur_prompt.replace("[FIELD]", ex.metadata.field)
        cur_prompt = cur_prompt.replace("[QUESTION]", ex.question)
        # resp = completion_with_backoff(
        #     model=FLAGS.model_name,
        #     prompt=cur_prompt,
        #     max_tokens=2048,
        #     )

        # chatgpt
        resp = chat_completion_with_backoff(
            model="gpt-4",
            messages=[{"role": "user", "content": cur_prompt}],
            max_tokens=2048,
        )
        output = resp["choices"][0]["message"]["content"]

        # output = resp["choices"][0]["text"]
        attribution_begin_idx = output.find("\n[1]")
        if attribution_begin_idx == -1:
            ex.answers[FLAGS.model_name] = example_utils.Answer(answer_string=output.strip())
        else:
            ex.answers[FLAGS.model_name] = example_utils.Answer(answer_string=output[:attribution_begin_idx].strip(),
                                                                  attribution=output[attribution_begin_idx+1:].split("\n"))
        json.dump(dataclasses.asdict(ex), f)
        f.write("\n")
    
    # example_utils.write_examples(FLAGS.output_file, input_data)


if __name__ == "__main__":
    app.run(main)

modeling/response_collection/fetch_openai_responses.py
chaitanyamalaviya/ExpertQA
import dataclasses
import json
import numpy as np
import os
import re
import requests
import sys
from typing import List

from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.schema import Document

from newspaper import Article
from newspaper.article import ArticleException

import openai
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)

from bs4 import BeautifulSoup # Required to parse HTML
from urllib.parse import unquote # Required to unquote URLs
import argparse

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

parser = argparse.ArgumentParser()
parser.add_argument("--input_file", help="input filepath", type=str)
parser.add_argument("--output_file", help="output filepath", type=str)
parser.add_argument("--start_idx", help="start index", type=int, default=0)
parser.add_argument("--topk", help="the value of k for the topk passages to use for QA", type=int, default=1)

args = parser.parse_args()


text_splitter = CharacterTextSplitter(separator=' ', chunk_size=1000, chunk_overlap=200)

#####
# OPENAI
#####
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
if OPENAI_API_KEY is None:
    raise ValueError("Please set the OpenAI API key as environment variable 'OPENAI_API_KEY'.")

openai.organization = ""
openai.api_key = ""

@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def chat_completion_with_backoff(**kwargs):
    return openai.ChatCompletion.create(**kwargs)


def google_search(
        query: str,
):
    """Get top 10 webpages from Google search.

    Args:
        query: search query
    Returns:
        top-10 search results in json format
    """
    response = requests.get(f"https://www.google.com/search?q={query}") # Make the request
    soup = BeautifulSoup(response.text, "html.parser") # Parse the HTML
    links = soup.find_all("a") # Find all the links in the HTML

    urls = []
    for l in [link for link in links if link["href"].startswith("/url?q=")]:
        # get the url
        url = l["href"]
        # remove the "/url?q=" part
        url = url.replace("/url?q=", "")
        # remove the part after the "&sa=..."
        url = unquote(url.split("&sa=")[0])
        # special case for google scholar
        if url.startswith("https://scholar.google.com/scholar_url?url=http"):
            url = url.replace("https://scholar.google.com/scholar_url?url=", "").split("&")[0]
        elif 'google.com/' in url: # skip google links
            continue
        if url.endswith('.pdf'): # skip pdf links
            continue
        if '#' in url: # remove anchors (e.g. wikipedia.com/bob#history and wikipedia.com/bob#genetics are the same page)
            url = url.split('#')[0]
        # print the url
        urls.append(url)

    # Use numpy to dedupe the list of urls after removing anchors
    urls = list(np.unique(urls))
    return urls


def scrape_and_parse(url: str):
    """Scrape a webpage and parse it into a Document object"""
    a = Article(url)
    try:
        a.download()
        a.parse()
    except Exception as e:
        return None

    return {
        "url": url,
        "text": a.text,
    }


def scrape_and_filter(urls: list):
    doc_list = []
    for u in urls:
        print(f"Processing: {u}")
        doc = scrape_and_parse(u)
        if doc is None:
            continue
        elif "Access" in doc["text"] and "Denied" in doc["text"]:
            continue
        else:
            doc_list.append(doc)

    return doc_list


def rerank(claim: str,
           retrieved_documents: List[Document],
           topk: int):

    texts = text_splitter.split_documents(retrieved_documents)
    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY, max_retries=1000)
    while True:
        try:
            docsearch = Chroma.from_documents(texts, embeddings)
        except Exception as e:
            continue
        break
    doc_retriever = docsearch.as_retriever(search_kwargs={"k": topk})
    topk_relevant_passages = doc_retriever.get_relevant_documents(claim)
    ret_passages = [{"text": d.page_content, "url": d.metadata["source"]} for d in topk_relevant_passages]
    return ret_passages


def remove_brackets(input_string):
    # Pattern to match substrings in square brackets
    pattern = r'\[.*?\]'
    # Substitute matching substrings with empty string
    output_string = re.sub(pattern, '', input_string)
    return output_string


if __name__ == '__main__':

    input_data = example_utils.read_examples(args.input_file)
    f = open(args.output_file, "a")
    for example in input_data[args.start_idx: args.start_idx + 200]:
        query_text = example.question
        post_hoc_claims = []
        for claim in example.answers["gpt4"]["claims"]:
            search_results = google_search(claim["claim_string"])

            if len(search_results) > 0:
                all_docs = scrape_and_filter(search_results)
            else:
                all_docs = []
                post_hoc_claims.append(example_utils.Claim(claim_string=claim["claim_string"], evidence=[]))
                continue                

            all_docs = [Document(page_content=d['text'], metadata={'source': d['url']}) for d in all_docs]
            all_docs_content_lens = [len(doc.page_content.strip()) for doc in all_docs]
            if not all_docs or not sum(all_docs_content_lens):
                post_hoc_claims.append(example_utils.Claim(claim_string=claim["claim_string"], evidence=[]))
                continue

            post_hoc_attributions = rerank(claim["claim_string"], all_docs, args.topk)
            post_hoc_claims.append(example_utils.Claim(claim_string=claim["claim_string"], evidence=post_hoc_attributions))
        example.answers["post_hoc_gs_gpt4"] = example.answers["gpt4"]
        example.answers["post_hoc_gs_gpt4"]["claims"] = post_hoc_claims
        example.answers["post_hoc_gs_gpt4"]["answer_string"] = remove_brackets(example.answers["gpt4"]["answer_string"])
        example.answers["post_hoc_gs_gpt4"]["attribution"] = []
        json.dump(dataclasses.asdict(example), f)
        f.write("\n")

modeling/response_collection/post_hoc_cite.py
chaitanyamalaviya/ExpertQA
import dataclasses
import json
import os
import re
import sys
from typing import List, Dict, Union

from pyserini.search import LuceneSearcher, FaissSearcher

import argparse

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

parser = argparse.ArgumentParser()
parser.add_argument("--input_file", help="input filepath", type=str)
parser.add_argument("--output_file", help="output filepath", type=str)
parser.add_argument("--ret_type", help="retriever type for wiki, dense or bm25",
                    type=str,
                    default="bm25",
                    choices=["dense", "bm25"])
parser.add_argument("--index_dir",
                    help="directory to processed wikipedia index",
                    default="/mnt/nlpgridio3/data/cmalaviya/sphere/faiss_index/sparse",
                    type=str)
parser.add_argument("--dense_index_model",
                    help="model used for indexing, only applies when using dense index.",
                    default="facebook/dpr-question_encoder-single-nq-base",
                    type=str)
parser.add_argument("--start_idx", help="start index", type=int, default=0)
parser.add_argument("--num_examples", help="Number of examples to process", type=int, default=200)
parser.add_argument("--topk", help="the value of k for the topk passages to use for QA", type=int, default=1)

args = parser.parse_args()


def sphere_search(query: str,
                  searcher: Union[LuceneSearcher, FaissSearcher],
                  top_k: int = 1):
    hits = searcher.search(query, k=top_k)
    retrieved = []
    for hit in hits:
        para_text = hit.raw
        retrieved.append({
            "url": json.loads(hit.docid)["url"],
            "text": para_text
        })

    return retrieved


def init_searcher(index_dir: str,
                  is_dense: bool = False,
                  dense_index_model: str = None):
    """
    Initialize the searcher
    """
    if is_dense:
        assert dense_index_model is not None, "with dense model, --dense_index_model needs to be set to the model that were used to create the index"
        print(f"Initializing the dense searcher from {index_dir} (model {dense_index_model})...")
        searcher = FaissSearcher(index_dir, dense_index_model)
    else:
        print(f"Initializing the BM25 searcher from {index_dir}...")
        searcher = LuceneSearcher(index_dir)

    return searcher


def remove_brackets(input_string):
    # Pattern to match substrings in square brackets
    pattern = r'\[.*?\]'
    # Substitute matching substrings with empty string
    output_string = re.sub(pattern, '', input_string)
    return output_string


if __name__ == '__main__':

    input_data = example_utils.read_examples(args.input_file)
    f = open(args.output_file, "w")

    if args.ret_type == "dense":
        searcher = init_searcher(index_dir=args.index_dir,
                                 is_dense=True,
                                 dense_index_model=args.dense_index_model)
    elif args.ret_type == "bm25":
        searcher = init_searcher(index_dir=args.index_dir,
                                 is_dense=False)
    else:
        raise ValueError(f"Retrieval type {args.wiki_ret_type} not supported.")

    for example in input_data[args.start_idx: args.start_idx + args.num_examples]:
        query_text = example.question
        post_hoc_claims = []
        for claim in example.answers["gpt4"]["claims"]:
            posthoc_attributions = sphere_search(
                query=claim["claim_string"],
                searcher=searcher,
                top_k=args.topk
            )

            post_hoc_claims.append(
                example_utils.Claim(claim_string=claim["claim_string"],
                                    evidence=posthoc_attributions)
            )

        example.answers[f"gpt4_post_hoc_sphere_{args.ret_type}"] = example.answers["gpt4"]
        example.answers[f"gpt4_post_hoc_sphere_{args.ret_type}"]["claims"] = post_hoc_claims
        example.answers[f"gpt4_post_hoc_sphere_{args.ret_type}"]["answer_string"] = remove_brackets(example.answers["gpt4"]["answer_string"])
        example.answers[f"gpt4_post_hoc_sphere_{args.ret_type}"]["attribution"] = []
        ex_json = json.dumps(dataclasses.asdict(example))
        f.write(ex_json)
        f.write("\n")

modeling/response_collection/post_hoc_cite_sphere.py
chaitanyamalaviya/ExpertQA
import dataclasses
import json
import numpy as np
import os
import requests
import sys
from typing import List

from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain import VectorDBQA
from langchain.chains import RetrievalQA, RetrievalQAWithSourcesChain
from langchain.schema import Document
from langchain.prompts import PromptTemplate

from newspaper import Article
from newspaper.article import ArticleException

import openai
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)

from bs4 import BeautifulSoup # Required to parse HTML
from urllib.parse import unquote # Required to unquote URLs


import argparse

prompt_template = """Use the following pieces of context to answer the question completely and precisely in up to 500 words. If you don't know the answer, just say "I don't know" and explain why the context is insufficient to answer the question.
You need to support every statement in the answer with in-line citations to passages given in the the context. The citations should appear as numbers such as [1], [2] that refer to the Passage IDs of the given passages. A statement may need to be supported by multiple references and should then be cited as [1] [2]. (for example, "Paris is the capital of France [1] [2]." where "1" and "2" are the Passage IDs of the first and second passage).

{context}

Question: {question}
Answer:"""

PROMPT = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

parser = argparse.ArgumentParser()
parser.add_argument("--input_file", help="input filepath", type=str)
parser.add_argument("--output_file", help="output filepath", type=str)
parser.add_argument("--topk", help="the value of k for the topk passages to use for QA", type=int, default=5)

args = parser.parse_args()


text_splitter = CharacterTextSplitter(separator=' ', chunk_size=1000, chunk_overlap=200)

#####
# OPENAI
#####
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
if OPENAI_API_KEY is None:
    raise ValueError("Please set the OpenAI API key as environment variable 'OPENAI_API_KEY'.")

openai.organization = ""
openai.api_key = ""

llm = ChatOpenAI(model_name='gpt-4', openai_api_key=OPENAI_API_KEY)

#####
# Google Search
#####
CSE_URL = "https://www.googleapis.com/customsearch/v1"
CSE_URL = "https://cse.google.com/cse?cx=74509b47ac2e54393"
gs_api_key = os.getenv('CUSTOM_SEARCH_API_KEY')
pse_cx = os.getenv('CUSTOM_SEARCH_CX')
if gs_api_key is None:
    raise ValueError("Please set the Custom search API key as environment variable 'CUSTOM_SEARCH_API_KEY'.")

if pse_cx is None:
    raise ValueError("Please set the Programmable search engine ID as environment variable 'CUSTOM_SEARCH_CX'.")

@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def chat_completion_with_backoff(**kwargs):
    return openai.ChatCompletion.create(**kwargs)


def google_search(
        query: str,
        api_key: str = None,
        cx: str = None
):
    """Get top 10 webpages from Google search.

    Args:
        query: search query
        api_key: custom search engine api key
        cx: programmable search engine id
    Returns:
        top-10 search results in json format
    """
    response = requests.get(f"https://www.google.com/search?q={query}") # Make the request
    soup = BeautifulSoup(response.text, "html.parser") # Parse the HTML
    links = soup.find_all("a") # Find all the links in the HTML

    urls = []
    for l in [link for link in links if link["href"].startswith("/url?q=")]:
        # get the url
        url = l["href"]
        # remove the "/url?q=" part
        url = url.replace("/url?q=", "")
        # remove the part after the "&sa=..."
        url = unquote(url.split("&sa=")[0])
        # special case for google scholar
        if url.startswith("https://scholar.google.com/scholar_url?url=http"):
            url = url.replace("https://scholar.google.com/scholar_url?url=", "").split("&")[0]
        elif 'google.com/' in url: # skip google links
            continue
        if url.endswith('.pdf'): # skip pdf links
            continue
        if '#' in url: # remove anchors (e.g. wikipedia.com/bob#history and wikipedia.com/bob#genetics are the same page)
            url = url.split('#')[0]
        # print the url
        urls.append(url)

    # Use numpy to dedupe the list of urls after removing anchors
    urls = list(np.unique(urls))
    return urls


    # if api_key is None:
    #     api_key = gs_api_key

    # if cx is None:
    #     cx = pse_cx

    # res = requests.get(
    #     url=CSE_URL,
    #     params={
    #         "q": query,
    #         "key": api_key,
    #         "cx": cx,
    #     },
    # )
    # if res.status_code != 200:
    #     print(f"Google search error: {res.status_code}")
    #     return []

    # res = res.json()

    # if 'items' in res:
    #     return res['items']
    # else:
    #     return []


def scrape_and_parse(url: str):
    """Scrape a webpage and parse it into a Document object"""
    a = Article(url)
    try:
        a.download()
        a.parse()
    except Exception as e:
        return None

    return {
        "url": url,
        "text": a.text,
    }


def scrape_and_filter(urls: list):
    doc_list = []
    for u in urls:
        print(f"Processing: {u}")
        doc = scrape_and_parse(u)
        if doc is None:
            continue
        elif "Access" in doc["text"] and "Denied" in doc["text"]:
            continue
        else:
            doc_list.append(doc)

    return doc_list


def retrieval_gpt_generate(query: str,
                           retrieved_documents: List[Document],
                           topk: int):

    texts = text_splitter.split_documents(retrieved_documents)
    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY, max_retries=1000)
    while True:
        try:
            docsearch = Chroma.from_documents(texts, embeddings)
        except Exception as e:
            continue
        break
    doc_retriever = docsearch.as_retriever(search_kwargs={"k": topk})
    topk_relevant_passages = doc_retriever.get_relevant_documents(query)
    topk_relevant_passages_content = ["Passage ID " + str(i+1) + ": " + doc.page_content.replace("\n","") for i, doc in enumerate(topk_relevant_passages)]
    ret_passages = [{"text": p, "url": d.metadata["source"]} for p, d in zip(topk_relevant_passages_content, topk_relevant_passages)]
    cur_prompt = prompt_template.format(context="Context: \n" + "\n\n".join(topk_relevant_passages_content), question=query)
    resp = chat_completion_with_backoff(
        model="gpt-4",
        messages=[{"role": "user", "content": cur_prompt}],
        max_tokens=2048,
    )
    answer = resp["choices"][0]["message"]["content"]
    # chain_type_kwargs = {"prompt": PROMPT}
    # qa = VectorDBQA.from_chain_type(llm=llm, chain_type="stuff", vectorstore=docsearch, return_source_documents=True)
    # qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=docsearch.as_retriever(search_kwargs={"k": topk}), return_source_documents=True, chain_type_kwargs=chain_type_kwargs)
    # result = qa({"query": query})

    return answer, ret_passages


if __name__ == '__main__':

    input_data = example_utils.read_examples(args.input_file)

    f = open(args.output_file, "a")
    for example in input_data:
        query_text = example.question

        # Google search first
        search_results = google_search(query_text)
        # search_urls = [r['link'] for r in search_results]
        if len(search_results) > 0:
            all_docs = scrape_and_filter(search_results)
        else:
            all_docs = []
            example.answers["ret_read_gpt4"] = example_utils.Answer(answer_string="I don't know.",
                                                                    attribution=[])
            json.dump(dataclasses.asdict(example), f)
            f.write("\n")
            continue

        all_docs = [Document(page_content=d['text'], metadata={'source': d['url']}) for d in all_docs]
        all_docs_content_lens = [len(doc.page_content.strip()) for doc in all_docs]
        if not all_docs or not sum(all_docs_content_lens):
            example.answers["ret_read_gpt4"] = example_utils.Answer(answer_string="I don't know.",
                                                                    attribution=[])
            json.dump(dataclasses.asdict(example), f)
            f.write("\n")
            continue
        gpt_query_text = "I am an expert in the field of " + example.metadata.field + ". Please answer my question: " + query_text
        answer, attributions = retrieval_gpt_generate(gpt_query_text, all_docs, args.topk)
        example.answers["rr_gs_gpt4"] = example_utils.Answer(answer_string=answer,
                                                                attribution=attributions)
        json.dump(dataclasses.asdict(example), f)
        f.write("\n")

modeling/response_collection/retrieve_and_read.py
chaitanyamalaviya/ExpertQA
import dataclasses
import json
import os
import sys
from typing import List, Dict, Union
from tqdm import tqdm

from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate

from pyserini.search import LuceneSearcher, FaissSearcher
import openai
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)

from data_utils import example_utils


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def chat_completion_with_backoff(**kwargs):
    return openai.ChatCompletion.create(**kwargs)


def sphere_search(query: str,
                  searcher: Union[LuceneSearcher, FaissSearcher],
                  top_k: int = 5):
    hits = searcher.search(query, k=top_k)
    retrieved = []
    for hit in hits:
        para_text = hit.raw
        retrieved.append({
            "url": json.loads(hit.docid)["url"],
            "text": para_text
        })

    return retrieved

def init_searcher(index_dir: str,
                  is_dense: bool = False,
                  dense_index_model: str = None):
    """
    Initialize the searcher
    """
    if is_dense:
        assert dense_index_model is not None, "with dense model, --dense_index_model needs to be set to the model that were used to create the index"
        print(f"Initializing the dense searcher from {index_dir} (model {dense_index_model})...")
        searcher = FaissSearcher(index_dir, dense_index_model)
    else:
        print(f"Initializing the BM25 searcher from {index_dir}...")
        searcher = LuceneSearcher(index_dir)

    return searcher


def retrieval_gpt_generate(query: str,
                           paragraphs: List[str],
                           prompt_template: str,
                           gpt_model_name: str = "gpt-4"):

    topk_relevant_passages_content = ["Passage ID " + str(i+1) + ": " + doc.replace("\n","") for i, doc in enumerate(paragraphs)]
    cur_prompt = prompt_template.format(context="Context: \n" + "\n".join(topk_relevant_passages_content),
                                        question=query)

    resp = chat_completion_with_backoff(
        model=gpt_model_name,
        messages=[{"role": "user", "content": cur_prompt}],
        max_tokens=2048,
    )

    answer = resp["choices"][0]["message"]["content"]
    return answer


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file",
                        help="input data to process",
                        type=str)
    parser.add_argument("--index_dir",
                        help="directory to processed sphere index",
                        default="/mnt/nlpgridio3/data/cmalaviya/sphere/faiss_index/sparse",
                        type=str)
    parser.add_argument("--dense_index_model",
                        help="model used for indexing, only applies when using dense index.",
                        default="facebook/dpr-question_encoder-single-nq-base",
                        type=str)
    parser.add_argument("--output_file",
                        help="output filepath",
                        type=str)
    parser.add_argument("--dense",
                        help="enable dense retrieval",
                        action="store_true")
    parser.add_argument("--topk",
                        help="the value of k for the topk passages to use for QA",
                        type=int,
                        default=5)
    parser.add_argument("--gpt_model_name",
                        help="Name of gpt model to use",
                        type=str,
                        default="gpt-4")
    parser.add_argument("--start_idx",
                        help="Start index",
                        type=int,
                        default="0")

    args = parser.parse_args()

    # set PYTHONPATH
    sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

    # Initialize LLM
    OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
    if OPENAI_API_KEY is None:
        raise ValueError("Please set the OpenAI API key as environment variable 'OPENAI_API_KEY'.")

    llm = ChatOpenAI(model_name=args.gpt_model_name, openai_api_key=OPENAI_API_KEY)

    # Prepare prompt
    prompt_template = """Use the following pieces of context to answer the question completely and precisely in up to 500 words. If you don't know the answer, just say "I don't know" and explain why the context is insufficient to answer the question.
    You need to support every statement in the answer with in-line citations to passages given in the the context. The citations should appear as numbers such as [1], [2] that refer to the Passage IDs of the given passages. A statement may need to be supported by multiple references and should then be cited as [1] [2]. (for example, "Paris is the capital of France [1] [2]." where "1" and "2" are the Passage IDs of the first and second passage).

    {context}

    Question: {question}
    Answer:"""

    PROMPT = PromptTemplate(
        template=prompt_template, input_variables=["context", "question"]
    )

    # Initialize retriever + docid to text mapping
    searcher = init_searcher(index_dir=args.index_dir,
                             is_dense=args.dense,
                             dense_index_model=args.dense_index_model)

    # Process the data
    input_data = example_utils.read_examples(args.input_file)
    with open(args.output_file, 'a') as fout:
        print(f"Using {args.gpt_model_name} to generate answers...")
        for example in tqdm(input_data[args.start_idx: args.start_idx+1000]):
            query_text = example.question
            relevant_paras = sphere_search(query_text,
                                           searcher=searcher,
                                           top_k=args.topk)
            relevant_paras_text = [p["text"] for p in relevant_paras]

            answer = retrieval_gpt_generate(query=query_text,
                                            paragraphs=relevant_paras_text,
                                            prompt_template=prompt_template,
                                            gpt_model_name=args.gpt_model_name)

            attributions = []
            for p in relevant_paras:
                attributions.append({
                    'url': p["url"],
                    'text': p["text"]
                })

            example.answers[f"rr_sphere_gpt4"] = example_utils.Answer(answer_string=answer,
                                                                      attribution=attributions)
            example_str = json.dumps(dataclasses.asdict(example))
            fout.write(example_str)
            fout.write("\n")

modeling/response_collection/sphere_and_read.py
chaitanyamalaviya/ExpertQA
"""Given an answer, split it into claims either using sentence boundaries or using GPT."""

import openai
import nltk
from tqdm import tqdm
from absl import app
from absl import flags
from tenacity import (
    retry,
    stop_after_attempt,
    wait_random_exponential,
)
import os
import re
import sys
import spacy
from spacy.lang.en import English

nlp = English()
nlp.add_pipe(nlp.create_pipe('sentencizer'))
# nlp.add_pipe("sentencizer")

# import config

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from data_utils import example_utils

openai.organization = ""
openai.api_key = ""

MIN_CLAIM_LENGTH=5

FLAGS = flags.FLAGS

flags.DEFINE_string('input_file', 'data/r1_data_answers.jsonl', 'Input filepath.')
flags.DEFINE_string('output_file', 'data/r1_data_answers_claims.jsonl', 'Output filepath.')
flags.DEFINE_string('granularity', 'sentence', 'Boundary for determining claims.')
flags.DEFINE_string('split_model_name', 'gpt4', 'Model name (required for sub-sentence claims).')
flags.DEFINE_string('prompt_config', 'prompts/claims_prompt.txt', 'Prompt path (required for sub-sentence claims).')
flags.DEFINE_integer('limit', None, 'Limit to the number of examples processed')

# Filter out any claims that contain the following substrings
filter_out_claims = set(["Hello",
                         "Hello, this is Bing.", 
                         "this is Bing",
                         "Thank you for your question"
                         "glad you're",
                         "glad you",
                         "glad to help you",
                         "glad to hear",
                         "Bing",
                         "Hope this helps",
                         "Hope this is helpful",
                         "I can provide more information based on your answer", 
                         "are you interested in?", 
                         "Are you looking for", 
                         "what aspect are you interested", 
                         "I found some information that might help you", 
                         "I found some information", 
                         "happy to help you", 
                         "I hope this information",
                         "I hope this helps",
                         "I hope this answer helps you",
                         "I can provide more specific information",
                         "I can provide more details",
                         "I searched for your question and found some results.", 
                         "I can help you",
                         "Thank you for your question",
                         "That's an interesting question",
                         "not sure if I understand your question",
                         "I hope this answers your question",
                         "I can provide you with more specific resources",
                         "I can try to answer your question",
                         "I can provide more information",
                         "I can search for more",
                         "I can answer your question",
                         "I can give you more",
                         "I can try to summarise",
                         "I searched for your question and found some information",
                         "I can answer your question",
                         "I can provide you",
                         "I can try to find",
                         "I cannot give you a definitive answer",
                         "I might be able to provide",
                         "please let me know"])
cite_pattern = r'\[\^.*?\^\]'


@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(6))
def completion_with_backoff(**kwargs):
    return openai.Completion.create(**kwargs)


def fix_sentence_boundaries(claim_sents):
    modified_sentences = []
    for i in range(len(claim_sents)):
        if i == len(claim_sents) - 1:
            modified_sentences.append(claim_sents[i])
            break
        
        next_sentence = claim_sents[i+1]
        bracketed_strings = ''
        while next_sentence.startswith('[') and ']' in next_sentence:
            end_index = next_sentence.index(']') + 1
            bracketed_strings += next_sentence[:end_index]
            next_sentence = next_sentence[end_index:].lstrip()  # remove leading whitespaces

        modified_sentences.append(claim_sents[i] + ' ' + bracketed_strings)
        claim_sents[i+1] = next_sentence
    return modified_sentences


def remove_passage_numbers_and_ids(string):
    pattern = r"Passage (Number|ID) \d+(?:[:])?"
    return re.sub(pattern, "", string)


def main(unused_argv):

    input_data = example_utils.read_examples(FLAGS.input_file)

    with open(FLAGS.prompt_config) as f:
        prompt = f.read()

    for i, ex in enumerate(tqdm(input_data)):
        if i==FLAGS.limit:
            break
        if "post_hoc_gs_gpt4" in FLAGS.input_file:
            del ex.answers["gpt4"]
            ex.answers["post_hoc_gs_gpt4"] = ex.answers.pop("gpt4_post_hoc_gs")
        elif "post_hoc_sphere_gpt4" in FLAGS.input_file:
            del ex.answers["gpt4"]
            ex.answers["post_hoc_sphere_gpt4"] = ex.answers.pop("gpt4_post_hoc_sphere_bm25")
        elif "rr_gs_gpt4" in FLAGS.input_file:
            ex.answers["rr_gs_gpt4"] = ex.answers.pop("ret_read_gpt4")
        for answer_model_name, answer in ex.answers.items():        
            if not answer["answer_string"].strip():
                continue
            cleaned_answer = re.sub(cite_pattern, '', answer["answer_string"]).strip()
            if "bingchat" in FLAGS.input_file:
                cleaned_answer = cleaned_answer.replace("]\n\nAre", "]. \n\nAre").replace("]\n\nWhat", "]. \n\nWhat").replace("]\n\nWhich", "]. \n\nWhich").replace("]\n\nHow", "]. \n\nHow").replace("]\n\nDoes", "]. \n\nDoes")
                cleaned_answer = cleaned_answer.replace("😊", "")
            ex.answers[answer_model_name]["answer_string"] = cleaned_answer
            # claim_sents = nltk.sent_tokenize(cleaned_answer.replace("\n", " "))
            doc = nlp(cleaned_answer.replace("\n", " "))
            claim_sents = [sent.text for sent in list(doc.sents)]

            # Citations in text-davinci-003 appear after the sentence boundary
            if answer_model_name == "text-davinci-003":
                claim_sents = fix_sentence_boundaries(claim_sents)

            if FLAGS.granularity == "sentence":
                if "post_hoc" not in answer_model_name:
                    ex.answers[answer_model_name]["claims"] = []
                    for claim in claim_sents:
                        # Check if claim needs to be filtered out
                        filt = False
                        for f_claim in filter_out_claims:
                            if f_claim in claim:
                                filt = True
                                break
                        # Exclude very short claims as they are likely a result of incorrect sentence tokenization
                        if len(claim) <= MIN_CLAIM_LENGTH:
                            filt = True
                        # Exclude questions from claims
                        if claim[-1].strip() == "?":
                            filt = True
                        
                        # If claim is being filtered, also remove it from the answer string
                        if filt:
                            ex.answers[answer_model_name]["answer_string"] = ex.answers[answer_model_name]["answer_string"].replace(claim.strip(), "")
                            continue
                        # Remove any substring enclosed by '[^' and '^]' as these indicate index of citation for BingChat
                        cleaned_claim = re.sub(cite_pattern, '', claim)
                        # Find all citation indexes
                        citation_matches = [i for i in re.findall(r'\[(.*?)\]', cleaned_claim) if i.isdigit()]
                        # Check if any attribution exists for the answer
                        if answer["attribution"] is not None and len(answer["attribution"]):
                            # Check if each attribution is a dict (true for `rr` systems with "text" and "url" keys)
                            if isinstance(answer["attribution"][0], dict):
                                evidence = []
                                for c_idx in citation_matches:
                                    # Invalid citation index
                                    if int(c_idx)-1 >= len(answer["attribution"]):
                                        continue
                                    evidence.append("[" + c_idx + "] " + answer["attribution"][int(c_idx)-1]["url"] + "\n\n" + remove_passage_numbers_and_ids(answer["attribution"][int(c_idx)-1]["text"]))
                                ex.answers[answer_model_name]["claims"].append(example_utils.Claim(cleaned_claim,
                                                                                                   evidence=evidence))
                            # true for bing_chat and gpt4
                            else:
                                ex.answers[answer_model_name]["claims"].append(example_utils.Claim(cleaned_claim,
                                                                                                   evidence=[answer["attribution"][int(c_idx)-1] for c_idx in citation_matches]))
                        else:
                            ex.answers[answer_model_name]["claims"].append(example_utils.Claim(cleaned_claim,
                                                                                               evidence=[]))
                    # Update attribution list for `rr` systems (initially a list of dicts to a list of strings)
                    if answer["attribution"] is not None and len(answer["attribution"]) and isinstance(answer["attribution"][0], dict):
                        ex.answers[answer_model_name]["attribution"] = ["[" + str(i+1) + "] " + attr["url"] for i, attr in enumerate(answer["attribution"])]
                else:
                    # post-hoc citation processing (already split into claims)
                    evidence = []
                    ex.answers[answer_model_name]["attribution"] = []
                    claim_idx = 1
                    claims = ex.answers[answer_model_name]["claims"][:]
                    ex.answers[answer_model_name]["claims"] = []                    
                    for claim in claims:
                        cleaned_claim = re.sub(r'\[\d+\]', '', claim["claim_string"])
                        if not cleaned_claim.strip() or len(cleaned_claim) <= MIN_CLAIM_LENGTH:
                            continue
                        claim["claim_string"] = cleaned_claim
                        evidence = []
                        if claim["evidence"]:
                            evidence = [f"[{claim_idx}] " + claim["evidence"][0]["url"] + "\n\n" + claim["evidence"][0]["text"]]
                            claim["claim_string"] = claim["claim_string"][:-1] + f"[{claim_idx}]" + claim["claim_string"][-1]
                            ex.answers[answer_model_name]["attribution"].append(f"[{claim_idx}] " + claim["evidence"][0]["url"])
                            claim_idx += 1
                        claim["evidence"] = evidence
                        ex.answers[answer_model_name]["claims"].append(claim)
                    ex.answers[answer_model_name]["answer_string"] = " ".join([claim["claim_string"] for claim in ex.answers[answer_model_name]["claims"]])

            else:
                for sent in claim_sents:
                    cur_prompt = prompt[:]
                    cur_prompt = cur_prompt.replace("[SENTENCE]", sent)
                    resp = completion_with_backoff(
                        model=FLAGS.split_model_name,
                        prompt=cur_prompt,
                        max_tokens=200,
                    )
                    output = resp["choices"][0]["text"]
                    claim_strings = output.split("- ")
                    ex.answers[answer_model_name]["claims"] += [example_utils.Claim(claim, 
                                                                                    evidence=[answer["attribution"][int(c_idx)-1] for c_idx in citation_matches]) for claim in claim_strings if claim.strip()]  
    
    example_utils.write_examples(FLAGS.output_file, input_data)


if __name__ == "__main__":
    app.run(main)

modeling/response_collection/split_ans_to_claims.py
chaitanyamalaviya/ExpertQA