Update inference.py

inference.py

@@ -17,29 +17,36 @@ model.to(device)
 import json
 
 def token2json(words, labels):
-    result = []
+    result = {}
     current_entity = None
-    
-    for token, label in zip(words, labels):
+    current_text = []
+
+    for word, label in zip(words, labels):
         if label.startswith("B-"):
             if current_entity:
-                result.append(current_entity)
-            current_entity = {"type": label[2:], "text": token}
+                result[current_entity] = " ".join(current_text).strip()
+                current_text = []
+            current_entity = label[2:].lower()
+            current_text = [word]
         elif label.startswith("I-"):
-            if current_entity and current_entity["type"] == label[2:]:
-                current_entity["text"] += " " + token
+            if current_entity == label[2:].lower():
+                current_text.append(word)
             else:
+                # Gestione di sequenze I- non precedute da B-
                 if current_entity:
-                    result.append(current_entity)
-                current_entity = {"type": label[2:], "text": token}
-        else:  # "O" label
+                    result[current_entity] = " ".join(current_text).strip()
+                current_entity = label[2:].lower()
+                current_text = [word]
+        else:  # Label "O"
             if current_entity:
-                result.append(current_entity)
+                result[current_entity] = " ".join(current_text).strip()
                 current_entity = None
-    
+                current_text = []
+
+    # Aggiunge l'ultima entità se presente
     if current_entity:
-        result.append(current_entity)
-    
+        result[current_entity] = " ".join(current_text).strip()
+
     return json.dumps(result, ensure_ascii=False, indent=2)