emanuelaboros commited on
Commit
50980d2
·
verified ·
1 Parent(s): e19b8b0

Update generic_ner.py

Browse files
Files changed (1) hide show
  1. generic_ner.py +3 -3
generic_ner.py CHANGED
@@ -238,11 +238,11 @@ class MultitaskTokenClassificationPipeline(Pipeline):
238
  confidence_scores[task] = F.softmax(logits, dim=-1).tolist()[0]
239
 
240
  entities = {}
241
- for task, preds in predictions.items():
242
  words_list, preds_list, confidence_list = realign(
243
  text_sentence,
244
- preds[0],
245
- confidence_scores[task][0],
246
  self.tokenizer,
247
  self.id2label[task],
248
  )
 
238
  confidence_scores[task] = F.softmax(logits, dim=-1).tolist()[0]
239
 
240
  entities = {}
241
+ for task in predictions.keys():
242
  words_list, preds_list, confidence_list = realign(
243
  text_sentence,
244
+ predictions[task],
245
+ confidence_scores[task],
246
  self.tokenizer,
247
  self.id2label[task],
248
  )