import os import json import unicodedata import zipfile 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 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) 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 CẤU HÌNH TOÀN CỤC & HỆ THỐNG QUẢN LÝ SESSION CACHE # ===================================================================== VOCAB_FILE = "vocab.txt" CONFIG_FILE = "config.json" BATCH_SIZE = 64 MAX_LENGTH = 512 if not all(os.path.exists(f) for f in [VOCAB_FILE, CONFIG_FILE]): raise FileNotFoundError("Thiếu file cấu hình (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()} # Lưu trữ các ONNX session đã khởi tạo để tránh load lại file nhiều lần _SESSIONS = {} def get_onnx_session(model_name): """Hàm nạp động và lưu cache cho từng loại model ONNX""" if model_name not in _SESSIONS: if not os.path.exists(model_name): raise FileNotFoundError(f"Không tìm thấy file mô hình '{model_name}' ở thư mục hiện tại!") 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 _SESSIONS[model_name] = ort.InferenceSession(model_name, sess_options, providers=["CPUExecutionProvider"]) return _SESSIONS[model_name] # ===================================================================== # HÀM XỬ LÝ INFERENCE DẠNG GENERATOR (YIELD REAL-TIME OUTPUT) # ===================================================================== def ner_inference(file_obj, selected_entities, min_count, model_choice): if file_obj is None: yield "❌ Thất bại", "Vui lòng tải lên một file dữ liệu dạng .txt hoặc .zip" return yield f"⏳ Đang nạp hoặc kiểm tra mô hình {model_choice}...", "Đang chuẩn bị dữ liệu văn bản..." try: session = get_onnx_session(model_choice) except Exception as e: yield "❌ Lỗi hệ thống", f"Không thể kích hoạt mô hình: {str(e)}" return lines = [] file_path = file_obj.name # TRƯỜNG HỢP 1: FILE ĐẦU VÀO LÀ FILE ZIP if file_path.lower().endswith('.zip'): try: with zipfile.ZipFile(file_path, 'r') as zip_ref: txt_files = [f for f in zip_ref.namelist() if f.lower().endswith('.txt') and not f.startswith('__MACOSX')] if not txt_files: yield "❌ Thất bại", "Không tìm thấy file .txt nào hợp lệ bên trong file ZIP." return for txt_file in txt_files: with zip_ref.open(txt_file) as f: content = f.read().decode('utf-8', errors='ignore') file_lines = [line.strip() for line in content.splitlines() if line.strip()] lines.extend(file_lines) except Exception as e: yield "❌ Lỗi đọc file ZIP", str(e) return # TRƯỜNG HỢP 2: FILE ĐẦU VÀO LÀ FILE TEXT ĐƠN LẺ else: try: with open(file_path, "r", encoding="utf-8", errors='ignore') as f: lines = [line.strip() for line in f if line.strip()] except Exception as e: yield "❌ Lỗi đọc file TXT", str(e) return if not lines: yield "❌ Thất bại", "Không trích xuất được dòng văn bản nào hợp lệ." return lines.sort(key=len) batches = [lines[i:i + BATCH_SIZE] for i in range(0, len(lines), BATCH_SIZE)] entity_count = defaultdict(int) total_batches = len(batches) # Vòng lặp tính toán và stream dữ liệu cuốn chiếu liên tục for idx, batch_lines in enumerate(batches): progress_msg = f"⏳ Đang quét: Batch {idx + 1}/{total_batches} (Tổng cộng {len(lines)} dòng văn bản)" 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"] } 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 # Sắp xếp danh sách thực thể theo số lần xuất hiện giảm dần để giao diện hiển thị tối ưu nhất sorted_entities = sorted(entity_count.items(), key=lambda x: (-x[1], x[0][0])) # Lọc và tạo chuỗi văn bản kết quả tính đến thời điểm hiện tại output_lines = [] for (name, label), count in sorted_entities: if count >= min_count and (not selected_entities or label in selected_entities): output_lines.append(f"{name}={label}={count}") current_stream_text = "\n".join(output_lines) if output_lines else "Chưa tìm thấy thực thể nào vượt qua ngưỡng min_count..." # Đẩy trạng thái tiến độ và danh sách cập nhật trực tiếp lên UI yield progress_msg, current_stream_text # Trả về kết quả chốt sau khi hoàn thành toàn bộ tác vụ final_progress = f"🎉 Hoàn thành xuất sắc! Đã quét xong toàn bộ {total_batches} batches." if not output_lines: yield final_progress, "Không tìm thấy thực thể nào khớp với bộ lọc hiện tại sau khi quét hết file." else: yield final_progress, "\n".join(output_lines) # ===================================================================== # THIẾT KẾ GIAO DIỆN GRADIO UI NÂNG CẤP # ===================================================================== css = "h1 { text-align: center; color: #1A202C; }" with gr.Blocks() as demo: gr.Markdown("# BERT NER Bulk Inference Engine (Real-time Stream Edition)") gr.Markdown("Hỗ trợ tải lên file đơn `.txt` hoặc tập hợp `.zip`. Hệ thống hiển thị tiến độ và cập nhật tần suất thực thể trực tiếp.") with gr.Row(): with gr.Column(scale=1): input_file = gr.File(label="Upload File (.txt hoặc .zip)", file_types=[".txt", ".zip"]) # Tính năng mới: Chọn mô hình chạy trực tiếp model_choice = gr.Dropdown( label="Chọn mô hình ONNX Inference", choices=["bert_ner_int8.onnx", "bert_ner_fp32.onnx"], value="bert_ner_int8.onnx" ) 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=500, step=1, value=1 ) submit_btn = gr.Button("Extract Entities", variant="primary") with gr.Column(scale=1): # Output 1: Hiển thị tiến trình thay đổi liên tục progress_text = gr.Textbox( label="Tiến độ thực thi (Progress Status)", interactive=False, lines=2, placeholder="Đang chờ lệnh từ người dùng..." ) # Output 2: Stream dữ liệu số lượng thực thể tăng tiến output_text = gr.Textbox( label="Kết quả thực thể Real-time (Stream Output)", interactive=False, lines=15, max_lines=25, placeholder="Danh sách thực thể kèm số lượng sẽ nhảy số liên tục tại đây..." ) submit_btn.click( fn=ner_inference, inputs=[input_file, entity_filter, count_entities, model_choice], outputs=[progress_text, output_text] ) demo.launch( css=css, theme=gr.themes.Soft() )