Update model.py

ChatWithPDF · Jan 22, 2024 · b6c1dfb · b6c1dfb
1 parent c942dd5
commit b6c1dfb
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/ner/agri_ner_akai/local/model.py b/src/ner/agri_ner_akai/local/model.py
@@ -44,10 +44,11 @@ def aggregate_entities(sentence, entity_outputs):
                     current_entity = None
 
                 else:
+                    if current_entity:
                     # If it's a subword but not an I- entity
-                    current_entity["word"] += entity["word"][2:]
-                    current_entity["end"] = entity["end"]
-                    current_entity["score"] = float((current_entity["score"] + entity["score"]) / 2)  # averaging scores
+                        current_entity["word"] += entity["word"][2:]
+                        current_entity["end"] = entity["end"]
+                        current_entity["score"] = float((current_entity["score"] + entity["score"]) / 2)  # averaging scores
 
             # Handle full words
             else: