Update inference.py

inference.py

@@ -14,6 +14,34 @@ label2id = model.config.label2id
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 
+import json
+
+def token2json(words, labels):
+    result = []
+    current_entity = None
+    
+    for token, label in zip(words, labels):
+        if label.startswith("B-"):
+            if current_entity:
+                result.append(current_entity)
+            current_entity = {"type": label[2:], "text": token}
+        elif label.startswith("I-"):
+            if current_entity and current_entity["type"] == label[2:]:
+                current_entity["text"] += " " + token
+            else:
+                if current_entity:
+                    result.append(current_entity)
+                current_entity = {"type": label[2:], "text": token}
+        else:  # "O" label
+            if current_entity:
+                result.append(current_entity)
+                current_entity = None
+    
+    if current_entity:
+        result.append(current_entity)
+    
+    return json.dumps(result, ensure_ascii=False, indent=2)
+
 
 def prediction(image):
     
@@ -55,18 +83,18 @@ def prediction(image):
 
     for i, conf in enumerate(true_confidence_scores):
         if conf < 0.6 :
-            true_predictions[i] = "O"          
+            true_predictions[i] = "O"        
+
     
-    d = {}
-    for id, i in enumerate(true_predictions):
+    d = token2json(true_words, true_predictions)
+    
+    """for id, i in enumerate(true_predictions):
         if i not in d.keys():
             d[i] = true_words[id]
         else:
             d[i] = d[i] + ", " + true_words[id]
     d = {k: v.strip() for (k, v) in d.items()}
-    d.pop("O")
-
-    # TODO:process the json
+    d.pop("O")"""
 
     draw = ImageDraw.Draw(image, "RGBA")
     font = ImageFont.load_default()