-
Notifications
You must be signed in to change notification settings - Fork 0
/
run.py
85 lines (77 loc) · 2.54 KB
/
run.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
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
from argparse import ArgumentParser
from brain2kg.text2kg.eda_pipeline import EDA
open('examples.log', 'w').close() # clear local log
if __name__ == '__main__':
parser = ArgumentParser()
# OIE module setting
parser.add_argument(
'--oie_llm',
default='llama3.1',
help='LLM used for open information extraction.',
)
parser.add_argument(
'--oie_prompt_template_file_path',
default='./prompt_templates/oie_fsp_template.txt',
help='Prompt template used for open information extraction.',
)
parser.add_argument(
'--oie_few_shot_example_file_path',
default='./few_shot_examples/oie_few_shot_examples.txt',
help='Few shot examples used for open information extraction.',
)
# SD module setting
parser.add_argument(
'--sd_llm',
default='llama3.1',
help='LLM used for schema definition.',
)
parser.add_argument(
'--sd_prompt_template_file_path',
default='./prompt_templates/sd_fsp_template.txt',
help='Prompt template used for schema definition.',
)
parser.add_argument(
'--sd_few_shot_example_file_path',
default='./few_shot_examples/sd_few_shot_examples.txt',
help='Few shot examples used for schema definition.',
)
# SA module setting
parser.add_argument(
'--sa_target_schema_file_path',
default='./schemas/webnlg_schema.csv',
help='Schema used for schema alignment verification.',
)
parser.add_argument(
'--sa_llm',
default='llama3.1',
help='LLM used for schema alignment verification.',
)
parser.add_argument(
'--sa_embedding_model',
default='sentence-transformers/all-MiniLM-L6-v2',
help='Embedding model used for schema alignment verification.',
)
parser.add_argument(
'--sa_prompt_template_file_path',
default='./prompt_templates/sa_template.txt',
help='Prompt template used for schema alignment verification.',
)
# Input text setting
parser.add_argument(
'--input_text_file_path',
default='./data/raw_text/webnlg_dataset.txt',
help='File containing input texts to extract KG from.',
)
# Output setting
parser.add_argument(
'--output_dir',
default='./examples/outputs',
help='Directory to output to.',
)
args = parser.parse_args()
args = vars(args)
eda = EDA(**args)
output_kg = eda.extract_kg(
args['input_text_file_path'],
args['output_dir']
)