| from argparse import ArgumentParser |
| from typing import * |
| import json |
| import logging |
|
|
| from sftp import SpanPredictor |
|
|
| logger = logging.getLogger('ConcretePredictor') |
|
|
|
|
| def read_kairos(ontology_mapping_path: Optional[str] = None): |
| |
| if ontology_mapping_path is None: |
| return |
| raw = json.load(open(ontology_mapping_path)) |
| fn2kairos = dict() |
| for kairos_label in raw: |
| for fn in raw[kairos_label]['framenet']: |
| fn_label = fn['label'] |
| if fn_label in fn2kairos: |
| logger.warning(f'"{fn_label}" is repeated in the ontology file.') |
| fn2kairos[fn_label] = kairos_label |
| return fn2kairos |
|
|
|
|
| def run(src, dst, model_path, ontology_mapping_path, device): |
| mapping = SpanPredictor.read_ontology_mapping(ontology_mapping_path) |
| predictor = SpanPredictor.from_path(model_path, cuda_device=device) |
| predictor.predict_concrete(src, dst, ontology_mapping=mapping) |
|
|
|
|
| if __name__ == '__main__': |
| parser = ArgumentParser() |
| parser.add_argument('src', type=str) |
| parser.add_argument('dst', type=str) |
| parser.add_argument('model', type=str) |
| parser.add_argument('--map', type=str, default=None) |
| parser.add_argument('--device', type=int, default=-1) |
| args = parser.parse_args() |
| run(args.src, args.dst, args.model, args.map, args.device) |
|
|