diff --git a/microservices/inferer/inferer.py b/microservices/inferer/inferer.py index 2853ba87ed635f04411b22485b9a500c2a42a584..1538c6530c8ed507d42b2117ac17cf5662a08e75 100644 --- a/microservices/inferer/inferer.py +++ b/microservices/inferer/inferer.py @@ -114,6 +114,13 @@ def inference_process(inference_data, model_id): attentions=outputs.attentions, ) + tag2id = {'action': 1, 'actor': 2, 'artifact': 3, 'condition': 4, 'location': 5, 'modality': 6, 'reference': 7, + 'time': 8} + label2id = { + 'O': 0, + **{f'{k}': v for k, v in tag2id.items()} + } + id2label = {v: k for k, v in label2id.items()} model = RobertaForSpanCategorization.from_pretrained(model_id) tokenizer = AutoTokenizer.from_pretrained(model_id)