preprocessing.py 22 KB
Newer Older
zihanl's avatar
zihanl committed
1

zihanl's avatar
zihanl committed
2
3
"""Preprocessing for Wizard of Wikipedia and Wizard of Internet datasets"""

root's avatar
root committed
4
import torch
zihanl's avatar
zihanl committed
5
6
7
import argparse
from nltk import word_tokenize
from tqdm import tqdm
zihanl's avatar
zihanl committed
8
9
import numpy as np
import json
zihanl's avatar
zihanl committed
10

root's avatar
root committed
11
def get_args():
zihanl's avatar
zihanl committed
12
13
    parser = argparse.ArgumentParser(description="Preprocessing")

root's avatar
root committed
14
    parser.add_argument("--func", type=str, default=None,
zihanl's avatar
zihanl committed
15
                        help="choose to run which function")
zihanl's avatar
zihanl committed
16
    parser.add_argument("--raw_file", type=str, default=None,
zihanl's avatar
zihanl committed
17
                        help="path of the input file")
zihanl's avatar
zihanl committed
18
19
20
21
22
23
24
25
    parser.add_argument("--processed_file", type=str, default=None,
                        help="path of the output file")
    parser.add_argument("--knwl_ref_file", type=str, default=None,
                        help="path of the knowledge reference file")
    parser.add_argument("--resp_ref_file", type=str, default=None,
                        help="path of the knowledge reference file")
    parser.add_argument("--knwl_gen_file", type=str, default=None,
                        help="path of the generated knowledge file")
root's avatar
root committed
26
    parser.add_argument("--test_file", type=str, default=None,
zihanl's avatar
zihanl committed
27
                        help="path of the test file")
root's avatar
root committed
28
    parser.add_argument("--train_file", type=str, default=None,
zihanl's avatar
zihanl committed
29
                        help="path of the train file")
root's avatar
root committed
30
    parser.add_argument("--model_file", type=str, default=None,
zihanl's avatar
zihanl committed
31
                        help="path of the model file")
root's avatar
root committed
32
    parser.add_argument("--data_type", type=str, default=None,
zihanl's avatar
zihanl committed
33
34
                        help="data types, choose one out of three types: \
                              wow_seen, wow_unseen, and woi")
root's avatar
root committed
35
    parser.add_argument("--seed", type=int, default=1234,
zihanl's avatar
zihanl committed
36
                        help="random seed")
zihanl's avatar
zihanl committed
37

root's avatar
root committed
38
39
    args = parser.parse_args()
    return args
zihanl's avatar
zihanl committed
40
41


zihanl's avatar
zihanl committed
42
def process_wow_dataset(raw_file, processed_file, knwl_ref_file, resp_ref_file):
zihanl's avatar
zihanl committed
43
    """
zihanl's avatar
zihanl committed
44
45
      This is a function used for processing the wizard of wikipedia (wow) dataset
      Expected processed format:
zihanl's avatar
zihanl committed
46
47
      topic \t dialogue context \t golden knowledge \t golden response
    """
zihanl's avatar
zihanl committed
48

zihanl's avatar
zihanl committed
49
50
    print("> Loading data from %s" % raw_file)
    with open(raw_file, "r") as fr:
zihanl's avatar
zihanl committed
51
52
        dialog_data = json.load(fr)
    
root's avatar
root committed
53
    print("> Processing data ...")
zihanl's avatar
zihanl committed
54
55
56
57
58
59
60
61
62
63
64
65
66
    fproc = open(processed_file, "w")
    fknwl = open(knwl_ref_file, "w") if knwl_ref_file else None
    fresp = open(resp_ref_file, "w") if resp_ref_file else None
    
    for i, sample in enumerate(tqdm(dialog_data)):
        # get all the dialog data for a single sample
        dialog = sample["dialog"]
        
        context = []
        for j, turn in enumerate(dialog):
            text = turn["text"]
            if not (text.endswith("?") or text.endswith(".") or text.endswith("!")):
                text = text + "."
zihanl's avatar
zihanl committed
67
            
zihanl's avatar
zihanl committed
68
69
70
71
72
73
74
75
76
            if j == 0:
                # first turn
                context.append(text)
                continue

            speaker = turn["speaker"].lower()
            if "wizard" in speaker:
                checked_sentence = list(turn["checked_sentence"].values())  # knowledge
                checked_passage = list(turn["checked_passage"].values())    # topic
zihanl's avatar
zihanl committed
77
                
zihanl's avatar
zihanl committed
78
                assert len(checked_sentence) <= 1
zihanl's avatar
zihanl committed
79

zihanl's avatar
zihanl committed
80
81
82
83
84
                # get the ground truth knowledge
                if len(checked_sentence) > 0:
                    checked_sentence = checked_sentence[0]
                else:
                    checked_sentence = "no_passages_used"
zihanl's avatar
zihanl committed
85

zihanl's avatar
zihanl committed
86
87
88
89
                if len(checked_passage) == 1:
                    checked_passage = checked_passage[0]
                else:
                    checked_passage = "no_passages_used"
zihanl's avatar
zihanl committed
90

zihanl's avatar
zihanl committed
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
                # get the topic
                if checked_passage != "no_passages_used":
                    topic = checked_passage
                else:
                    topic = sample["chosen_topic"]
                
                knowledge = checked_sentence
                response = text
                # write to the output files
                fproc.write(topic + "\t" + " [SEP] ".join(context) + "\t" + \
                                knowledge + "\t" + response + "\n")
                
                if fknwl:
                    fknwl.write(knowledge + "\n")
                if fresp:
                    # tokenize for evaluation
                    response = " ".join(word_tokenize(response))
                    fresp.write(response + "\n")
zihanl's avatar
zihanl committed
109

zihanl's avatar
zihanl committed
110
                context.append(text)
zihanl's avatar
zihanl committed
111

zihanl's avatar
zihanl committed
112
113
114
115
116
117
118
119
120
            else:
                assert "apprentice" in speaker
                context.append(text)

    fproc.close()
    if fknwl:
        fknwl.close()
    if fresp:
        fresp.close()
zihanl's avatar
zihanl committed
121
122


zihanl's avatar
zihanl committed
123
def process_woi_dataset(raw_file, processed_file, knwl_ref_file, resp_ref_file):
zihanl's avatar
zihanl committed
124
125
126
127
128
    """
      This is a function used for processing the wizard of internet (woi) dataset
      Expected processed format:
      topic \t dialogue context \t golden knowledge \t golden response
    """
zihanl's avatar
zihanl committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
    
    print("> Processing %s" % raw_file)
    fproc = open(processed_file, "w")
    fknwl = open(knwl_ref_file, "w") if knwl_ref_file else None
    fresp = open(resp_ref_file, "w") if resp_ref_file else None
    
    with open(raw_file, "r") as fr:
        for i, line in tqdm(enumerate(fr)):
            line = line.strip()
            item_dict = json.loads(line)
            item_dict = item_dict.values()
            assert len(item_dict) == 1
            item_dict = list(item_dict)[0]
            
            dialog_data = item_dict['dialog_history']
            length = len(dialog_data)
            
            turn_list = []
            search_text = ""
            for i in range(length):
                item = dialog_data[i]
                action = item['action']
zihanl's avatar
zihanl committed
151

zihanl's avatar
zihanl committed
152
153
154
155
156
                if action == "Wizard => SearchAgent":
                    search_text = item['text']

                elif action == "Wizard => Apprentice":
                    if len(turn_list) == 0:
zihanl's avatar
zihanl committed
157
158
                        turn = item['text']
                        turn_list.append(turn)
zihanl's avatar
zihanl committed
159
160
161
162
163
164
165
166
167
168
169
170
171
172
                        continue

                    # get the relevant content
                    contents = item["context"]["contents"]
                    selects = item["context"]["selected_contents"]
                    flag = selects[0][0]
                    selects = selects[1:]
                    assert len(selects) == len(contents)
                    
                    # get the topic
                    if flag:
                        # no knowledge sentence is used
                        topic = "no_topic"
                        sent_list = ["no_passages_used"]
zihanl's avatar
zihanl committed
173
                    else:
zihanl's avatar
zihanl committed
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
                        # assert search_text != ""
                        topic = search_text

                        sent_list = []
                        for content, select in zip(contents, selects):
                            content = content['content']
                            assert len(content) == len(select)
                            for c, s in zip(content, select):
                                if s:
                                    sent_list.append(c)
                    if len(sent_list) == 0:
                        topic = "no_topic"
                        sent_list = ["no_passages_used"]

                    # get dialogue context, knowledge, and response 
                    dialog_context = " [SEP] ".join(turn_list)
                    knwl_sent = sent_list[0]
                    response = item['text']

                    # processing
                    topic = topic.replace("\n", "").replace("\r", \
                                "").replace("\t", "")
                    dialog_context = dialog_context.replace("\n", "").replace("\r", \
                                "").replace("\t", "")
                    knwl_sent = knwl_sent.replace("\n", "").replace("\r", \
                                "").replace("\t", "")
                    response = response.replace("\n", "").replace("\r", \
                                "").replace("\t", "")
                    
                    if topic != "no_topic":
                        # write to the ouput files
                        fproc.write(topic + "\t" + dialog_context + "\t" + \
                                        knwl_sent + "\t" + response + "\n")
                        if fknwl:
                            fknwl.write(knwl_sent + "\n")
                        if fresp:
                            # tokenize for evaluation
                            response = " ".join(word_tokenize(response))
                            fresp.write(response + "\n")

                    turn_list.append(response)

                elif action == "Apprentice => Wizard":
                    turn = item['text']
                    turn_list.append(turn)

                else:
                    assert action == "SearchAgent => Wizard"

    fproc.close()
    if fknwl:
        fknwl.close()
    if fresp:
        fresp.close()
zihanl's avatar
zihanl committed
228
229


root's avatar
root committed
230
231
def get_database(test_datapath, train_datapath, data_type):
    """Get the database by topics"""
zihanl's avatar
zihanl committed
232

zihanl's avatar
zihanl committed
233
234
    assert data_type in ["wow_seen", "wow_unseen", "woi"], \
                "Please input a correct data type!!"
zihanl's avatar
zihanl committed
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
    # get test data topic list
    print("> reading test data from %s" % test_datapath)
    test_topics = {}
    with open(test_datapath, "r") as f:
        for i, line in enumerate(f):
            line = line.strip()
            splits = line.split("\t")
            topic = splits[0]
            test_topics[topic] = True

    print("> reading data from %s" % train_datapath)
    train_data_by_topic = {}
    dialog_data_by_topic = {}
    dialog_examples = []
    with open(train_datapath, "r") as f:
        for i, line in enumerate(f):
            line = line.strip()
            splits = line.split("\t")
            topic = splits[0]
            turns = splits[1].split(" [SEP] ")[-3:]
            knowledge = splits[2]
            response = splits[3]
root's avatar
root committed
257
            # filtering data samples
zihanl's avatar
zihanl committed
258
259
            if knowledge == "no_passages_used":
                continue
root's avatar
root committed
260
261
262
263
264
            if data_type != "wow_seen" and ("(" in knowledge or ")" in knowledge):
                continue
            if data_type != "wow_seen" and topic not in knowledge:
                continue

zihanl's avatar
zihanl committed
265
266
            # get the instance
            last_turn = turns[-1]
root's avatar
root committed
267
268
269
270
            if data_type == "woi":
                instance = "( " + last_turn + " ) " + topic + " -> " + knowledge
            else:
                instance = "( " + last_turn + " ) " + topic + " => " + knowledge
zihanl's avatar
zihanl committed
271
272
273
            
            # construct dialog example
            dialog_example = ""
root's avatar
root committed
274
275
276
277
278
            if data_type != "wow_seen":
                dialog_example += "( " + topic + " ) "
            for i, turn in enumerate(turns):
                if i != 0:
                    dialog_example += " "
zihanl's avatar
zihanl committed
279
                dialog_example += turn
root's avatar
root committed
280
            
zihanl's avatar
zihanl committed
281
282
283
284
285
286
287
288
289
290
291
            # check overlaps
            if topic in test_topics:
                if topic not in train_data_by_topic:
                    train_data_by_topic[topic] = [instance]
                else:
                    train_data_by_topic[topic].append(instance)
                
                if topic not in dialog_data_by_topic:
                    dialog_data_by_topic[topic] = [dialog_example]
                else:
                    dialog_data_by_topic[topic].append(dialog_example)
root's avatar
root committed
292
293
294
295
296
297
298
299
300
301
            
            else:
                # filtering data samples
                if len(knowledge.split()) > 20:
                    # knowledge is too long
                    continue
                if knowledge.startswith("It") or knowledge.startswith("it") or \
                   knowledge.startswith("This") or knowledge.startswith("this"):
                    continue
                
zihanl's avatar
zihanl committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
            # append all the data into dialogue examples list
            dialog_examples.append((topic, dialog_example, instance))

    return train_data_by_topic, dialog_data_by_topic, dialog_examples


emb_dict = {}
def select_prompts_based_on_similarity(
        query, dialog_list, prompt_list, topic, tokenizer, encoder, topk):
    """Select samples based on the similarity"""

    with torch.no_grad():
        # get the query embeddings
        query_ids = tokenizer.encode(query)
        query_ids = torch.LongTensor([query_ids]).cuda()
        query_emb = encoder(input_ids=query_ids).pooler_output
        query_emb = query_emb[0]
        
        # calculate embeddings for the samples in the database
        if topic in emb_dict:
            example_embeddings = emb_dict[topic]
            example_embeddings = example_embeddings.cuda()
        else:
            for idx, example in enumerate(dialog_list):
                example_ids = tokenizer.encode(example)
                example_ids = torch.LongTensor([example_ids]).cuda()
                example_emb = encoder(input_ids=example_ids).pooler_output
                if idx == 0:
                    example_embeddings = example_emb
                else:
                    example_embeddings = torch.cat(
                        (example_embeddings, example_emb), dim=0)
            emb_dict[topic] = example_embeddings.cpu()

        # compare the similarity and select the topk samples
        similarity_list = example_embeddings.matmul(query_emb)
        _, indices = torch.topk(similarity_list, k=topk)
    
    indices = indices.tolist()
    indices = indices[::-1] # reverse the order
    selected_prompts = []
    for index in indices:
        # index = index.item()
        selected_prompts.append(prompt_list[index])

    return selected_prompts


def prompt_selection_for_knowledge_generation(
root's avatar
root committed
351
        test_datapath, train_datapath, model_path, output_prompt_path, data_type):
zihanl's avatar
zihanl committed
352
353
354
355
356
    """Selecting prompts for the knowledge generation"""

    print("> Selecting prompts for the knowledge generation")

    train_data_by_topic, dialog_data_by_topic, dialog_examples = \
root's avatar
root committed
357
                            get_database(test_datapath, train_datapath, data_type)
zihanl's avatar
zihanl committed
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
    
    from transformers import DPRQuestionEncoderTokenizer
    print("> loading tokenizer and encoder")
    tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
                    'facebook/dpr-question_encoder-single-nq-base')
    encoder = torch.load(model_path).cuda()

    print("> getting dialog embeddings")
    with torch.no_grad():
        for idx, example in tqdm(enumerate(dialog_examples)):
            dialog = example[1]
            dialog_ids = tokenizer.encode(dialog)
            dialog_ids = torch.LongTensor([dialog_ids]).cuda()
            dialog_emb = encoder(input_ids=dialog_ids).pooler_output

            if idx == 0:
                dialog_embeddings = dialog_emb
            else:
                dialog_embeddings = torch.cat((dialog_embeddings, dialog_emb), dim=0)

    print("> reading test data from %s" % test_datapath)
    prompt_list_for_each_sample = []
    with open(test_datapath, "r") as f:
        for i, line in tqdm(enumerate(f)):
            line = line.strip()

            splits = line.split("\t")
            topic = splits[0]
            turns = splits[1].split(" [SEP] ")[-3:]

root's avatar
root committed
388
389
390
391
392
393
394
395
            # get the query sentence
            query_sent = ""
            if data_type != "seen":
                query_sent += "( " + topic + " ) "
            for i, turn in enumerate(turns):
                if i != 0:
                    query_sent += " "
                query_sent += turn
zihanl's avatar
zihanl committed
396

root's avatar
root committed
397
            if topic not in train_data_by_topic:
zihanl's avatar
zihanl committed
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
                # get the query embedding
                query_ids = tokenizer.encode(query_sent)
                query_ids = torch.LongTensor([query_ids]).cuda()
                query_emb = encoder(input_ids=query_ids).pooler_output
                query_emb = query_emb[0]

                # calculate the similarity
                similarity_list = dialog_embeddings.matmul(query_emb)
                _, indices = torch.sort(similarity_list)
                indices = indices.tolist()
                selected_topics = {}
                selected_prompts = []
                num_prompt = 0
                for index in indices:
                    example = dialog_examples[index]
                    topic_temp = example[0]
                    if topic_temp not in selected_topics:
                        selected_topics[topic_temp] = True
                        selected_prompts.append(example[2])
                        num_prompt += 1
                        if num_prompt == 10:
                            break
                
                # get the selected samples
                example_list = selected_prompts[::-1]
                key = topic + " " + turns[-1]
                prompt_list_for_each_sample.append({key: example_list})

            else:
                num_data_sample = min(len(train_data_by_topic[topic]), 10)
                total_example_list = train_data_by_topic[topic]
root's avatar
root committed
429
                
zihanl's avatar
zihanl committed
430
                dialog_list = dialog_data_by_topic[topic]
root's avatar
root committed
431
                assert len(dialog_list) == len(train_data_by_topic[topic])
zihanl's avatar
zihanl committed
432
433

                # calculate the similarity
root's avatar
root committed
434
                example_list = select_prompts_based_on_similarity(
zihanl's avatar
zihanl committed
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
                                query_sent, dialog_list, total_example_list, 
                                topic, tokenizer, encoder, topk=num_data_sample)
                
                key = topic + " " + turns[-1]
                prompt_list_for_each_sample.append({key: example_list})

    print("writing to %s" % output_prompt_path)
    with open(output_prompt_path, "w") as f:
        for instance in tqdm(prompt_list_for_each_sample):
            json.dump(instance, f)
            f.write("\n")


def prompt_selection_for_response_generation(input_path, output_path, seed):
    """Selecting prompts for the response generation"""

    print("> Selecting prompts for the response generation")
    print("> set random seed")
    np.random.seed(seed)

    prompt_example_list = []
    print("> reading data from %s" % input_path)
    with open(input_path, "r") as f:
        for i, line in tqdm(enumerate(f)):
            line = line.strip()
            splits = line.split("\t")

            # get the topic, context, knowledge and response
            topic = splits[0]
            dialog_context = splits[1]
            knowledge = splits[2]
            response = splits[3]
            turns = dialog_context.split(" [SEP] ")[-3:]
            if knowledge == "no_passages_used":
                continue

            # calculate the overlap ratio
            from nltk import word_tokenize
            knowledge_sent_token_list = word_tokenize(knowledge)
            knowledge_sent_token_dict = {token: True for token in knowledge_sent_token_list}
root's avatar
root committed
475
476
            knowledge_len = len(knowledge_sent_token_list)
            response_token_list = word_tokenize(response)
zihanl's avatar
zihanl committed
477
478
            response_len = len(response_token_list)
            num_overlap_token = 0
root's avatar
root committed
479
            accumulator = 0
zihanl's avatar
zihanl committed
480
481
            for token in response_token_list:
                if token in knowledge_sent_token_dict:
root's avatar
root committed
482
483
484
485
486
487
488
                    accumulator += 1
                else:
                    if accumulator >= 10:
                        num_overlap_token += accumulator
                    accumulator = 0
            if accumulator >= 10:
                num_overlap_token += accumulator
zihanl's avatar
zihanl committed
489
490
491
492
            
            # filtering the data based on the ratio
            if num_overlap_token > response_len * 0.9 or num_overlap_token < response_len * 0.6:
                continue
root's avatar
root committed
493
494
495
496
497
498
            if num_overlap_token < knowledge_len * 0.8:
                continue
            
            last_turn = " ".join(word_tokenize(turns[-1]))
            knowledge = " ".join(word_tokenize(knowledge))
            response = " ".join(word_tokenize(response))
zihanl's avatar
zihanl committed
499
500
501
            prompt_example = ""
            # add dialog context
            prompt_example += "Topic: " + topic + ". "
root's avatar
root committed
502
            prompt_example += "User says: " + last_turn + " "
zihanl's avatar
zihanl committed
503
504
505
506
507
            prompt_example += "We know that: " + knowledge + " "
            prompt_example += "System replies: " + response
            
            prompt_example_list.append(prompt_example)
        
root's avatar
root committed
508
    # shuffle the prompt examples
zihanl's avatar
zihanl committed
509
510
511
512
513
514
515
516
517
518
    np.random.shuffle(prompt_example_list)
    
    print("> writing to %s" % output_path)
    with open(output_path, "w") as f:
        # f.write("Generate the System's response based on the knowledge sentence:\n")
        for i in tqdm(range(20)):
            example = prompt_example_list[i]
            f.write(example + "\n")


zihanl's avatar
zihanl committed
519
def prepare_input_for_response_generation(test_file, knwl_gen_file, processed_file):
zihanl's avatar
zihanl committed
520
521
    """Preparing inputs for the response generation"""

zihanl's avatar
zihanl committed
522
    print("> Reading knowledge file from %s" % knwl_gen_file)
zihanl's avatar
zihanl committed
523
    # get the knowledge list
zihanl's avatar
zihanl committed
524
    with open(knwl_gen_file, "r") as f:
zihanl's avatar
zihanl committed
525
526
        knowledge_list = f.readlines()
    
root's avatar
root committed
527
    print("> Processing ...")
zihanl's avatar
zihanl committed
528
    with open(test_file, "r") as fr:
zihanl's avatar
zihanl committed
529
        with open(processed_file, "w") as fw:
zihanl's avatar
zihanl committed
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
            for line_num, line in enumerate(tqdm(fr)):
                line = line.strip()
                splits = line.split("\t")
                # prepare topic, context, knowledge and response
                topic = splits[0]
                dialog_context = splits[1]
                response = splits[3]
                knowledge = knowledge_list[line_num]
                knowledge = knowledge.strip()
                if "<|endoftext|>" in knowledge:
                    knowledge = knowledge.replace("<|endoftext|>", "")

                # write to the output file
                fw.write(topic + "\t" + dialog_context + "\t" \
                                     + knowledge + "\t" + response + "\n")

zihanl's avatar
zihanl committed
546
547
548

if __name__ == "__main__":

root's avatar
root committed
549
550
    args = get_args()
    if args.func == "process_wow_dataset":
zihanl's avatar
zihanl committed
551
        process_wow_dataset(args.raw_file, args.processed_file, args.knwl_ref_file, args.resp_ref_file)
zihanl's avatar
zihanl committed
552

root's avatar
root committed
553
    elif args.func == "process_woi_dataset":
zihanl's avatar
zihanl committed
554
        process_woi_dataset(args.raw_file, args.processed_file, args.knwl_ref_file, args.resp_ref_file)
zihanl's avatar
zihanl committed
555

root's avatar
root committed
556
    elif args.func == "get_knwl_gen_prompts":
zihanl's avatar
zihanl committed
557
        prompt_selection_for_knowledge_generation(
root's avatar
root committed
558
            args.test_file, args.train_file, args.model_file, 
zihanl's avatar
zihanl committed
559
            args.processed_file, args.data_type)
root's avatar
root committed
560
561
    
    elif args.func == "get_resp_gen_prompts":
zihanl's avatar
zihanl committed
562
        prompt_selection_for_response_generation(
zihanl's avatar
zihanl committed
563
            args.train_file, args.processed_file, args.seed)
zihanl's avatar
zihanl committed
564

root's avatar
root committed
565
    elif args.func == "prepare_input":
zihanl's avatar
zihanl committed
566
        prepare_input_for_response_generation(
zihanl's avatar
zihanl committed
567
            args.test_file, args.knwl_gen_file, args.processed_file)