fix extraction

main.py

@@ -247,7 +247,8 @@ async def extract_entities(input_data: TextInput):
     predictions = outputs.logits.argmax(dim=2)
     entities = {}
     current_entity = None
-    current_text = []
+    current_start = None
+    current_end = None
 
     word_ids = inputs.word_ids(0)
     for idx, word_idx in enumerate(word_ids):
@@ -260,31 +261,35 @@ async def extract_entities(input_data: TextInput):
         predicted_label = id2label[prediction]
 
         if predicted_label.startswith("B-"):
-            if current_entity:
+            if current_entity is not None:
                 entity_type = current_entity[2:]
                 if entity_type not in entities:
-                    entities[entity_type] = [" ".join(current_text)]
-                # Only keep the first detection, do nothing if already present
+                    entities[entity_type] = [text[current_start:current_end]]
+                current_entity = None
+                current_start = None
+                current_end = None
 
             current_entity = predicted_label
-            current_text = [clean_tokens[word_idx]]
+            current_start, current_end = token_positions[word_idx]
 
         elif predicted_label.startswith("I-") and current_entity and predicted_label[2:] == current_entity[2:]:
-            current_text.append(clean_tokens[word_idx])
+            # Extend the end position to include this token
+            _, token_end = token_positions[word_idx]
+            current_end = token_end
 
         else:
-            if current_entity:
+            if current_entity is not None:
                 entity_type = current_entity[2:]
                 if entity_type not in entities:
-                    entities[entity_type] = [" ".join(current_text)]
-                # Only keep the first detection, do nothing if already present
+                    entities[entity_type] = [text[current_start:current_end]]
                 current_entity = None
-                current_text = []
+                current_start = None
+                current_end = None
 
-    if current_entity:
+    if current_entity is not None:
         entity_type = current_entity[2:]
         if entity_type not in entities:
-            entities[entity_type] = [" ".join(current_text)]
+            entities[entity_type] = [text[current_start:current_end]]
         # Only keep the first detection, do nothing if already present
 
     return {"entities": entities}