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,15 +12,27 @@ 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) -> Dict[str, Any]:
+def ner(input: str) -> str:
     output = get_completion(input)
     merged_tokens = merge_tokens(output)
-    return {"text": input, "entities": merged_tokens}
+    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)
 
 css = '''
 h1#title {