Update app.py

app.py

@@ -1,46 +1,361 @@
-import gradio as gr
-import requests
+import os
 import json
+import unicodedata
+from collections import defaultdict
+import numpy as np
+import onnxruntime as ort
+import gradio as gr
 
+# =====================================================================
+# CHUẨN ĐẶC TẢ BERT TOKENIZER (KEEP ORIGINAL LOGIC)
+# =====================================================================
+class PureBertTokenizer:
+    def __init__(self, vocab_file, do_lower_case=None):
+        self.vocab = {}
+        self.id_to_vocab = {}
+        with open(vocab_file, "r", encoding="utf-8") as f:
+            for idx, line in enumerate(f):
+                token = line.strip("\r\n")
+                self.vocab[token] = idx
+                self.id_to_vocab[idx] = token
+        
+        if do_lower_case is None:
+            has_uppercase = any(any(c.isupper() for c in t) for t in self.vocab if not t.startswith("["))
+            self.do_lower_case = not has_uppercase
+        else:
+            self.do_lower_case = do_lower_case
 
-url = 'https://api.coze.com/open_api/v2/chat'
-headers = {
-    'Authorization': 'Bearer pat_6d5X5P8ykxPvtmOzaetgOAWIn4YI5lX99YdUWwcEDT241YZYunmJQLo9Y5oToEpb',
-    'Content-Type': 'application/json',
-    'Accept': '*/*',
-    'Host': 'api.coze.com',
-    'Connection': 'keep-alive',
-}
-def print_like_dislike(x: gr.LikeData):
-    print(x.index, x.value, x.liked)
-    
+    def _clean_text(self, text):
+        output = []
+        for char in text:
+            cp = ord(char)
+            if cp == 0 or cp == 0xfffd or self._is_control(char):
+                continue
+            if self._is_whitespace(char):
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
 
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(
-        elem_id="Xombot",
-        bubble_full_width=False,
-        label = "Xombot"
-    )
-    msg = gr.Textbox(interactive=True, placeholder="Đặt câu hỏi ở đây...", show_label=False)
-    clear = gr.ClearButton([msg, chatbot])
-    
-    def respond(query, history):
-        data = {
-        "conversation_id": "123",
-        "bot_id": "7369918746766868488",
-        "user": "996999",
-        "query": f"{query}",
-        "stream": False
+    def _is_whitespace(self, char):
+        if char in [" ", "\t", "\n", "\r"]:
+            return True
+        return unicodedata.category(char) == "Zs"
+
+    def _is_control(self, char):
+        if char in [" ", "\t", "\n", "\r"]:
+            return False
+        return unicodedata.category(char).startswith("C")
+
+    def _tokenize_chinese_chars(self, text):
+        output = []
+        for char in text:
+            if self._is_chinese_char(ord(char)):
+                output.append(" ")
+                output.append(char)
+                output.append(" ")
+            else:
+                output.append(char)
+        return "".join(output)
+
+    def _is_chinese_char(self, cp):
+        if ((cp >= 0x4E00 and cp <= 0x9FFF) or (cp >= 0x3400 and cp <= 0x4DBF) or
+            (cp >= 0x20000 and cp <= 0x2A6DF) or (cp >= 0x2A700 and cp <= 0x2B73F) or
+            (cp >= 0x2B740 and cp <= 0x2B81F) or (cp >= 0x2B820 and cp <= 0x2CEAF) or
+            (cp >= 0xF900 and cp <= 0xFAFF) or (cp >= 0x2F800 and cp <= 0x2FA1F)):
+            return True
+        return False
+
+    def _run_strip_accents(self, text):
+        text = unicodedata.normalize("NFD", text)
+        output = []
+        for char in text:
+            if unicodedata.category(char) == "Mn":
+                continue
+            output.append(char)
+        return "".join(output)
+
+    def _run_split_on_punc(self, text):
+        chars = list(text)
+        i = 0
+        start_new_token = True
+        output = []
+        while i < len(chars):
+            char = chars[i]
+            if self._is_punctuation(char):
+                output.append([char])
+                start_new_token = True
+            else:
+                if start_new_token:
+                    output.append([])
+                start_new_token = False
+                output[-1].append(char)
+            i += 1
+        return ["".join(x) for x in output]
+
+    def _is_punctuation(self, char):
+        cp = ord(char)
+        if (33 <= cp <= 47) or (58 <= cp <= 64) or (91 <= cp <= 96) or (123 <= cp <= 126):
+            return True
+        return unicodedata.category(char).startswith("P")
+
+    def tokenize(self, text):
+        text = unicodedata.normalize("NFC", text)
+        text = self._clean_text(text)
+        text = self._tokenize_chinese_chars(text)
+        
+        orig_tokens = text.split()
+        split_tokens = []
+        for token in orig_tokens:
+            if self.do_lower_case:
+                token = token.lower()
+                token = self._run_strip_accents(token)
+            split_tokens.extend(self._run_split_on_punc(token))
+
+        output_tokens = []
+        for token in split_tokens:
+            chars = list(token)
+            if len(chars) > 100:
+                output_tokens.append("[UNK]")
+                continue
+
+            is_bad = False
+            start = 0
+            sub_tokens = []
+            while start < len(chars):
+                end = len(chars)
+                cur_substr = None
+                while start < end:
+                    substr = "".join(chars[start:end])
+                    if start > 0:
+                        substr = "##" + substr
+                    if substr in self.vocab:
+                        cur_substr = substr
+                        break
+                    end -= 1
+                if cur_substr is None:
+                    is_bad = True
+                    break
+                sub_tokens.append(cur_substr)
+                start = end
+
+            if is_bad:
+                output_tokens.append("[UNK]")
+            else:
+                output_tokens.extend(sub_tokens)
+        
+        return output_tokens
+
+    def encode_batch(self, texts, max_length=512, padding=True, truncation=True):
+        batch_input_ids, batch_attention_mask, batch_token_type_ids = [], [], []
+        max_len_in_batch = 0
+        tokenized_batch = []
+        
+        for text in texts:
+            tokens = self.tokenize(text)
+            if truncation and len(tokens) > max_length - 2:
+                tokens = tokens[:max_length - 2]
+            
+            input_ids = (
+                [self.vocab["[CLS]"]] + 
+                [self.vocab.get(t, self.vocab["[UNK]"]) for t in tokens] + 
+                [self.vocab["[SEP]"]]
+            )
+            tokenized_batch.append(input_ids)
+            if len(input_ids) > max_len_in_batch:
+                max_len_in_batch = len(input_ids)
+                
+        target_len = max_length if (padding and max_len_in_batch > max_length) else max_len_in_batch
+        if not padding:
+            target_len = max_len_in_batch
+            
+        for input_ids in tokenized_batch:
+            if len(input_ids) > target_len:
+                input_ids = input_ids[:target_len]
+            pad_len = target_len - len(input_ids)
+            
+            attention_mask = [1] * len(input_ids) + [0] * pad_len
+            token_type_ids = [0] * target_len
+            input_ids = input_ids + [self.vocab["[PAD]"]] * pad_len
+            
+            batch_input_ids.append(input_ids)
+            batch_attention_mask.append(attention_mask)
+            batch_token_type_ids.append(token_type_ids)
+            
+        return {
+            "input_ids": np.array(batch_input_ids, dtype=np.int64),
+            "attention_mask": np.array(batch_attention_mask, dtype=np.int64),
+            "token_type_ids": np.array(batch_token_type_ids, dtype=np.int64)
+        }
+        
+    def convert_ids_to_tokens(self, ids):
+        return [self.id_to_vocab.get(int(i), "[UNK]") for i in ids]
+
+
+def join_bert_tokens(token_list):
+    text = ""
+    for token in token_list:
+        if token.startswith("##"):
+            text += token[2:]
+        else:
+            if text and ('\u4e00' <= token <= '\u9fff' or ('\u4e00' <= text[-1] <= '\u9fff')):
+                text += token
+            else:
+                text += (" " if text else "") + token
+    return text.strip()
+
+
+# =====================================================================
+# KHỞI TẠO MÔ HÌNH TOÀN CỤC (LOAD ONCE)
+# =====================================================================
+ONNX_MODEL_PATH = "bert_ner_fp32.onnx"
+VOCAB_FILE = "vocab.txt"
+CONFIG_FILE = "config.json"
+BATCH_SIZE = 64          
+MAX_LENGTH = 512         
+
+# Kiểm tra file cấu hình bắt buộc trước khi load ứng dụng
+if not all(os.path.exists(f) for f in [ONNX_MODEL_PATH, VOCAB_FILE, CONFIG_FILE]):
+    raise FileNotFoundError("Thiếu file bert_ner_fp32.onnx, vocab.txt hoặc config.json ở thư mục hiện tại!")
+
+tokenizer = PureBertTokenizer(VOCAB_FILE)
+
+with open(CONFIG_FILE, "r", encoding="utf-8") as f:
+    config_data = json.load(f)
+id2label = {int(k): v for k, v in config_data["id2label"].items()}
+
+sess_options = ort.SessionOptions()
+sess_options.intra_op_num_threads = 0  
+sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+session = ort.InferenceSession(ONNX_MODEL_PATH, sess_options, providers=["CPUExecutionProvider"])
+
+
+# =====================================================================
+# HÀM XỬ LÝ GRADIO INFERENCE
+# =====================================================================
+def ner_inference(file_obj, selected_entities, min_count):
+    if file_obj is None:
+        return "Vui lòng tải lên một file dữ liệu dạng văn bản (.txt)"
+        
+    # Đọc dữ liệu đầu vào từ file tạm do Gradio tạo ra
+    with open(file_obj.name, "r", encoding="utf-8") as f:
+        lines = [line.strip() for line in f if line.strip()]
+
+    if not lines:
+        return "File được tải lên không có dữ liệu văn bản hợp lệ."
+
+    lines.sort(key=len)
+    batches = [lines[i:i + BATCH_SIZE] for i in range(0, len(lines), BATCH_SIZE)]
+    entity_count = defaultdict(int)
+
+    # Vòng lặp Inference từng batch
+    for batch_lines in batches:
+        encoded = tokenizer.encode_batch(batch_lines, padding=True, truncation=True, max_length=MAX_LENGTH)
+        
+        ort_inputs = {
+            "input_ids": encoded["input_ids"],
+            "attention_mask": encoded["attention_mask"],
+            "token_type_ids": encoded["token_type_ids"]
         }
-        response = requests.post(url, headers=headers, data=json.dumps(data))
-        response_json = response.json()
-        for message in response_json['messages']:
-            if message['role'] == 'assistant' and message['type'] == 'answer':
-                history.append((query, message['content']))
-                return "", history # message['content']
-              
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    chatbot.like(print_like_dislike, None, None)
+        
+        ort_outputs = session.run(["logits"], ort_inputs)
+        logits = ort_outputs[0]  
+        predictions = np.argmax(logits, axis=-1)
+        
+        for i in range(len(batch_lines)):
+            input_ids_seq = encoded["input_ids"][i]
+            attention_mask_seq = encoded["attention_mask"][i]
+            pred_seq = predictions[i]
+            
+            tokens = tokenizer.convert_ids_to_tokens(input_ids_seq)
+            current_entity_tokens = []
+            current_entity_label = None
+            
+            for token, mask, label_id in zip(tokens, attention_mask_seq, pred_seq):
+                if mask == 0 or token == "[SEP]":  
+                    break
+                if token == "[CLS]":
+                    continue
+                    
+                label = id2label[label_id]
+                prefix = label.split("-")[0] if "-" in label else label
+                ent_type = label.split("-")[1] if "-" in label else None
+                
+                if prefix in ["B", "S", "O"] or (prefix in ["M", "I", "E"] and ent_type != current_entity_label):
+                    current_entity_tokens = []
+                    current_entity_label = None
+                
+                if prefix == "S":
+                    final_name = join_bert_tokens([token])
+                    if final_name:
+                        entity_count[(final_name, ent_type)] += 1
+                elif prefix == "B":
+                    current_entity_label = ent_type
+                    current_entity_tokens.append(token)
+                elif prefix in ["M", "I"]:
+                    if current_entity_label == ent_type:
+                        current_entity_tokens.append(token)
+                elif prefix == "E":
+                    if current_entity_label == ent_type:
+                        current_entity_tokens.append(token)
+                        final_name = join_bert_tokens(current_entity_tokens)
+                        if final_name:
+                            entity_count[(final_name, current_entity_label)] += 1
+                    current_entity_tokens = []
+                    current_entity_label = None
+
+    # Lọc thực thể dựa trên cài đặt UI
+    output_lines = []
+    for (name, label), count in entity_count.items():
+        if count >= min_count and (not selected_entities or label in selected_entities):
+            output_lines.append(f"{name}={label}={count}")
+
+    if not output_lines:
+        return "Không tìm thấy thực thể nào tương ứng với các bộ lọc hiện tại."
+
+    return "\n".join(output_lines)
+
+
+# =====================================================================
+# THIẾT KẾ GIAO DIỆN GRADIO
+# =====================================================================
+css = "h1 { text-align: center; color: #2D3748; }"
+
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# BERT NER Inference Engine (ONNX Pure Python)")
+    gr.Markdown("Tải lên file văn bản `.txt` để trích xuất thực thể tên riêng theo mong muốn bằng mô hình ONNX.")
     
-if __name__ == "__main__":
-    demo.launch()
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_file = gr.File(label="Upload File (.txt)", file_types=[".txt"])
+            entity_filter = gr.CheckboxGroup(
+                label="Entities Filter",
+                choices=["PER", "ORG", "LOC", "GPE"],
+                value=["PER", "ORG", "LOC", "GPE"]
+            )
+            count_entities = gr.Number(
+                label="Min Frequency Threshold",
+                minimum=1,
+                maximum=50,
+                step=1,
+                value=1
+            )
+            submit_btn = gr.Button("Extract Entities", variant="primary")
+            
+        with gr.Column(scale=1):
+            output_text = gr.Textbox(
+                label="Output Results",
+                show_copy_button=True,
+                interactive=False,
+                lines=15,
+                max_lines=25
+            )
+            
+    submit_btn.click(
+        fn=ner_inference,
+        inputs=[input_file, entity_filter, count_entities],
+        outputs=[output_text]
+    )
+
+demo.launch()