From da9db8836a96cd78931f4a1b8f558ca6ed0e5704 Mon Sep 17 00:00:00 2001
From: "Julien B." <xm9q8f80@jlnbrtn.me>
Date: Sun, 25 Aug 2024 22:20:31 +0200
Subject: [PATCH] fix(inferer): fix missing var
---
microservices/inferer/inferer.py | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/microservices/inferer/inferer.py b/microservices/inferer/inferer.py
index 1538c65..a91290d 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]
--
GitLab