-
Notifications
You must be signed in to change notification settings - Fork 0
/
cleanup.py
26 lines (20 loc) · 697 Bytes
/
cleanup.py
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
import argparse
from mas import CnnDmWriter, XSumWriter, DUC2004Writer
DATA_WRITERS = {
'cnndm': CnnDmWriter,
'xsum': XSumWriter,
'duc2004': DUC2004Writer
}
def create_writer(args):
return DATA_WRITERS[args.config](args)
def main(args):
writer = create_writer(args)
writer.write()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, choices=['cnndm', 'xsum', 'duc2004'], required=True)
parser.add_argument('--input-dir', type=str, required=True)
parser.add_argument('--output-dir', type=str, required=True)
parser.add_argument('--max-len', type=int)
args = parser.parse_args()
main(args)