diff --git a/microservices/inferer/inferer.py b/microservices/inferer/inferer.py index 1538c6530c8ed507d42b2117ac17cf5662a08e75..a91290d3c6c746545ddb17be40ed4c84e9358921 100644 --- a/microservices/inferer/inferer.py +++ b/microservices/inferer/inferer.py @@ -125,9 +125,9 @@ def inference_process(inference_data, model_id): model = RobertaForSpanCategorization.from_pretrained(model_id) tokenizer = AutoTokenizer.from_pretrained(model_id) - def get_offsets_and_predicted_tags(example: str, model, tokenizer, threshold=0): - raw_encoded_example = tokenizer(example, return_offsets_mapping=True) - encoded_example = tokenizer(example, return_tensors="pt") + def get_offsets_and_predicted_tags(sentence: str, model, tokenizer, threshold=0): + raw_encoded_example = tokenizer(sentence, return_offsets_mapping=True) + encoded_example = tokenizer(sentence, return_tensors="pt") out = model(**encoded_example)["logits"][0] predicted_tags = [[i for i, l in enumerate(logit) if l > threshold] for logit in out]