app.py

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()
)