Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import pipeline
 from typing import List, Dict, Any
 
-def merge_tokens(tokens: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def merge_tokens(tokens: List[Dict[str, any]]) -> List[Dict[str, any]]:
     merged_tokens = []
     for token in tokens:
         if merged_tokens and token['entity'].startswith('I-') and merged_tokens[-1]['entity'].endswith(token['entity'][2:]):
@@ -12,27 +12,15 @@ def merge_tokens(tokens: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
             last_token['score'] = (last_token['score'] + token['score']) / 2
         else:
             merged_tokens.append(token)
+
     return merged_tokens
 
 get_completion = pipeline("ner", model="b3x0m/bert-xomlac-ner")
 
-def ner(input: str) -> str:
+def ner(input: str) -> Dict[str, Any]:
     output = get_completion(input)
     merged_tokens = merge_tokens(output)
-    entity_map = {
-        "PER": "tên người",
-        "LOC": "địa điểm",
-        "ORG": "tổ chức",
-        "MISC": "vị trí",
-    }
-    result = []
-    for token in merged_tokens:
-        entity = token['entity']
-        if entity in entity_map:  # Filter only relevant entities
-            entity_label = entity_map.get(entity, "khác")
-            result.append(f"{token['word']} ({entity_label})")
-    
-    return ", ".join(result)
+    return {"text": input, "entities": merged_tokens}
 
 css = '''
 h1#title {