File size: 5,988 Bytes
2b35800
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
147d3e2
2b35800
 
 
 
 
147d3e2
2b35800
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ed1749d
2b35800
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17da9c4
2b35800
 
 
 
 
 
 
17da9c4
 
 
2b35800
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
17da9c4
2b35800
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
import argparse
import time
import json
import nltk
from rank_bm25 import BM25Okapi
import numpy as np
import torch
from transformers import BloomTokenizerFast, BloomForCausalLM


def claim2prompts(example):
    claim = example["claim"]

    # claim_str = "Claim: " + claim + "||Evidence: "
    claim_str = "Evidence: "

    for question in example["questions"]:
        q_text = question["question"].strip()
        if len(q_text) == 0:
            continue

        if not q_text[-1] == "?":
            q_text += "?"

        answer_strings = []

        for a in question["answers"]:
            if a["answer_type"] in ["Extractive", "Abstractive"]:
                answer_strings.append(a["answer"])
            if a["answer_type"] == "Boolean":
                answer_strings.append(
                    a["answer"]
                    + ", because "
                    + a["boolean_explanation"].lower().strip()
                )

        for a_text in answer_strings:
            if not a_text[-1] in [".", "!", ":", "?"]:
                a_text += "."

            # prompt_lookup_str = claim + " " + a_text
            prompt_lookup_str = a_text
            this_q_claim_str = (
                claim_str + " " + a_text.strip() + "||Question answered: " + q_text
            )
            yield (
                prompt_lookup_str,
                this_q_claim_str.replace("\n", " ").replace("||", "\n"),
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Use a prompt to generate questions that could be answered by top-k retrieved evidence. Output generated questions."
    )
    parser.add_argument("--reference_corpus", default="data/train.json", help="")
    parser.add_argument("--target_file", default="data/dev.json", help="")
    parser.add_argument(
        "-i",
        "--top_k_target_knowledge",
        default="data_store/dev_top_k_sentences.json",
        help="Directory where the sentences for the scraped data is saved.",
    )
    parser.add_argument(
        "-o",
        "--output_questions",
        default="data_store/dev_top_k_qa.json",
        help="Directory where the sentences for the scraped data is saved.",
    )
    parser.add_argument(
        "--top_k",
        default=100,
        type=int,
        help="How many documents should we pick out with BM25",
    )
    args = parser.parse_args()

    # few-shot learning from the training set
    with open(args.reference_corpus, "r", encoding="utf-8") as json_file:
        train_examples = json.load(json_file)

    prompt_corpus, tokenized_corpus = [], []

    for example in train_examples:
        for lookup_str, prompt in claim2prompts(example):
            entry = nltk.word_tokenize(lookup_str)
            tokenized_corpus.append(entry)
            prompt_corpus.append(prompt)

    prompt_bm25 = BM25Okapi(tokenized_corpus)

    # Load the bloom model:
    tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
    model = BloomForCausalLM.from_pretrained(
        "bigscience/bloom-7b1",
        device_map="auto",
        torch_dtype=torch.bfloat16,
        offload_folder="./offload",
    )

    with open(args.output_questions, "w", encoding="utf-8") as output_file:
        with open(args.top_k_target_knowledge, "r", encoding="utf-8") as json_file:
            for i, line in enumerate(json_file):
                data = json.loads(line)
                top_k_sentences_urls = data[f"top_{args.top_k}"]
                claim = data["claim"]
                claim_id = data["claim_id"]

                bm25_qau = []  # question, answer, url
                # Generate questions for those top k:
                for sent_i, sentences_urls in enumerate(top_k_sentences_urls):

                    prompt_lookup_str = sentences_urls["sentence"]
                    url = sentences_urls["url"]

                    prompt_s = prompt_bm25.get_scores(
                        nltk.word_tokenize(prompt_lookup_str)
                    )
                    prompt_n = 10
                    prompt_top_n = np.argsort(prompt_s)[::-1][:prompt_n]
                    prompt_docs = [prompt_corpus[i] for i in prompt_top_n]

                    claim_prompt = (
                        "Evidence: "
                        + prompt_lookup_str.replace("\n", " ")
                        + "\nQuestion answered: "
                    )

                    prompt = "\n\n".join(prompt_docs + [claim_prompt])

                    inputs = tokenizer([prompt], padding=True, return_tensors="pt").to(
                        model.device
                    )
                    st = time.time()
                    outputs = model.generate(
                        inputs["input_ids"],
                        max_length=5000,
                        num_beams=2,
                        no_repeat_ngram_size=2,
                        early_stopping=True,
                    )
                    print(
                        f"Generated QA for sent {sent_i} in file {i}. Time elapsed: {time.time() - st}"
                    )

                    tgt_text = tokenizer.batch_decode(
                        outputs[:, inputs["input_ids"].shape[-1] :],
                        skip_special_tokens=True,
                    )[0]

                    # We are not allowed to generate more than 250 characters:
                    tgt_text = tgt_text[:250]

                    qau_pair = [
                        tgt_text.strip().split("?")[0].replace("\n", " ") + "?",
                        prompt_lookup_str.replace("\n", " "),
                        url,
                    ]

                    bm25_qau.append(qau_pair)

                json_data = {
                    "claim_id": claim_id,
                    "claim": claim,
                    "bm25_qau": bm25_qau,
                }
                output_file.write(json.dumps(json_data, ensure_ascii=False) + "\n")
                output_file.flush()