| """ |
| AI Translation System với Knowledge Graph + NER + Transliteration |
| ================================================================== |
| Pipeline 4 lớp xử lý tên riêng: |
| Lớp 1: Lookup Knowledge Graph DB (entity đã biết, độ chính xác cao nhất) |
| Lớp 2: NER tự động phát hiện tên chưa có trong DB |
| Lớp 3: Transliteration (unidecode) cho PERSON/LOCATION |
| Lớp 4: NMT Model dịch phần còn lại (tên đã được mask) |
| """ |
|
|
| import gradio as gr |
| import os |
| import re |
| import logging |
| from dataclasses import dataclass, field |
| from functools import lru_cache |
|
|
| import httpx |
| from unidecode import unidecode |
|
|
| |
| |
| |
| logging.basicConfig( |
| level=logging.INFO, |
| format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", |
| ) |
| logger = logging.getLogger("translation_app") |
|
|
|
|
| |
| |
| |
| @dataclass |
| class AppConfig: |
| supabase_url: str = field( |
| default_factory=lambda: os.environ.get( |
| "SUPABASE_URL", |
| "https://jkaazcjahfqozwizhpvk.supabase.co" |
| ) |
| ) |
| supabase_key: str = field( |
| default_factory=lambda: os.environ.get("SUPABASE_KEY", "") |
| ) |
| allowed_types: tuple = ( |
| "PERSON", "ORG", "LOCATION", |
| "TECH", "PROGRAMMING_LANGUAGE", "FRAMEWORK", |
| "DATABASE", "CLOUD", "MODEL", |
| "PRODUCT", "FOOD", |
| ) |
| min_frequency: int = 1 |
| |
| |
| ev_model: str = "LLC536/nmt-en-vi-translation-v1" |
| ve_model: str = "LLC536/nmt-vi-en-translation-v1" |
| |
|
|
| CONFIG = AppConfig() |
|
|
|
|
| |
| |
| |
| _models: dict = {} |
|
|
|
|
| def get_model(direction: str): |
| if direction not in _models: |
| from transformers import pipeline |
| model_id = CONFIG.ev_model if direction == "en->vi" else CONFIG.ve_model |
| logger.info(f"Đang load model dịch: {model_id}") |
| _models[direction] = pipeline("translation", model=model_id) |
| logger.info(f"Model [{direction}] sẵn sàng.") |
| return _models[direction] |
|
|
|
|
| |
| |
| |
| _ner_vi = None |
| _ner_en = None |
|
|
|
|
| def get_ner_vi(): |
| """NER tiếng Việt dùng underthesea.""" |
| global _ner_vi |
| if _ner_vi is None: |
| try: |
| from underthesea import ner as _ner |
| _ner_vi = _ner |
| logger.info("NER tiếng Việt (underthesea) sẵn sàng.") |
| except ImportError: |
| logger.warning("underthesea chưa cài. NER tiếng Việt bị tắt.") |
| _ner_vi = None |
| return _ner_vi |
|
|
|
|
| def get_ner_en(): |
| """NER tiếng Anh dùng transformers pipeline.""" |
| global _ner_en |
| if _ner_en is None: |
| try: |
| from transformers import pipeline |
| _ner_en = pipeline( |
| "ner", |
| model="dslim/bert-base-NER", |
| aggregation_strategy="simple", |
| ) |
| logger.info("NER tiếng Anh (dslim/bert-base-NER) sẵn sàng.") |
| except Exception as e: |
| logger.warning(f"Không load được NER tiếng Anh: {e}") |
| _ner_en = None |
| return _ner_en |
|
|
|
|
| |
| |
| |
| @lru_cache(maxsize=8) |
| def fetch_entities_from_supabase(lang_code: str) -> dict: |
| """ |
| Lấy entity glossary từ Supabase theo ngôn ngữ. |
| Cache lru_cache – chỉ gọi API 1 lần mỗi lang_code. |
| """ |
| if not CONFIG.supabase_key: |
| logger.warning("SUPABASE_KEY chưa cấu hình. KG bị bỏ qua.") |
| return {} |
|
|
| headers = { |
| "apikey": CONFIG.supabase_key, |
| "Authorization": f"Bearer {CONFIG.supabase_key}", |
| } |
| params = { |
| "select": "entity_name,normalize_name", |
| "language": f"eq.{lang_code}", |
| "frequency": f"gte.{CONFIG.min_frequency}", |
| "entity_type": f"in.({','.join(CONFIG.allowed_types)})", |
| } |
|
|
| try: |
| resp = httpx.get( |
| f"{CONFIG.supabase_url}/rest/v1/knowledge_graph_entities", |
| headers=headers, |
| params=params, |
| timeout=10.0, |
| ) |
| resp.raise_for_status() |
| entities = {} |
| for row in resp.json(): |
| name = (row.get("entity_name") or "").strip() |
| norm = (row.get("normalize_name") or "").strip() |
| if name and norm: |
| entities[name] = norm |
| logger.info(f"Loaded {len(entities)} entities từ Supabase (lang={lang_code}).") |
| return entities |
| except httpx.HTTPStatusError as e: |
| logger.error(f"Supabase HTTP {e.response.status_code}: {e.response.text}") |
| return {} |
| except Exception as e: |
| logger.error(f"Lỗi Supabase API: {e}", exc_info=True) |
| return {} |
|
|
|
|
| def invalidate_cache(): |
| fetch_entities_from_supabase.cache_clear() |
| logger.info("Entity cache đã xóa.") |
|
|
|
|
| |
| |
| |
| def _convert_by_type(name: str, entity_type: str) -> str: |
| """ |
| Quyết định cách xử lý tên theo loại thực thể: |
| - PERSON / LOC → bỏ dấu bằng unidecode |
| - ORG / TECH → giữ nguyên (tên tổ chức/kỹ thuật không dịch) |
| """ |
| if entity_type in ("PER", "PERSON", "LOC", "LOCATION"): |
| return unidecode(name) |
| return name |
|
|
|
|
| def auto_detect_entities_vi(text: str) -> dict: |
| """ |
| Dùng underthesea NER để phát hiện tên riêng trong tiếng Việt. |
| Trả về {tên_gốc: tên_đích}. |
| """ |
| ner_fn = get_ner_vi() |
| if not ner_fn: |
| return {} |
|
|
| try: |
| results = ner_fn(text) |
| |
| |
| detected = {} |
| current_words = [] |
| current_type = None |
|
|
| for item in results: |
| |
| word = item[0] |
| tag = item[-1] |
|
|
| if tag.startswith("B-"): |
| |
| if current_words and current_type: |
| name = " ".join(current_words) |
| detected[name] = _convert_by_type(name, current_type) |
| current_words = [word] |
| current_type = tag[2:] |
|
|
| elif tag.startswith("I-") and current_words: |
| current_words.append(word) |
|
|
| else: |
| |
| if current_words and current_type: |
| name = " ".join(current_words) |
| detected[name] = _convert_by_type(name, current_type) |
| current_words = [] |
| current_type = None |
|
|
| |
| if current_words and current_type: |
| name = " ".join(current_words) |
| detected[name] = _convert_by_type(name, current_type) |
|
|
| logger.info(f"NER VI phát hiện {len(detected)} entity: {list(detected.keys())}") |
| return detected |
|
|
| except Exception as e: |
| logger.error(f"Lỗi NER tiếng Việt: {e}", exc_info=True) |
| return {} |
|
|
|
|
| def auto_detect_entities_en(text: str) -> dict: |
| """ |
| Dùng bert-base-NER để phát hiện tên riêng trong tiếng Anh. |
| Trả về {tên_gốc: tên_đích} (giữ nguyên – EN→VI thì model lo). |
| """ |
| ner_fn = get_ner_en() |
| if not ner_fn: |
| return {} |
|
|
| try: |
| results = ner_fn(text) |
| detected = {} |
| for ent in results: |
| word = ent.get("word", "").strip() |
| ent_type = ent.get("entity_group", "") |
| score = ent.get("score", 0) |
| if word and score > 0.85: |
| |
| detected[word] = word |
| logger.info(f"NER EN phát hiện {len(detected)} entity: {list(detected.keys())}") |
| return detected |
| except Exception as e: |
| logger.error(f"Lỗi NER tiếng Anh: {e}", exc_info=True) |
| return {} |
|
|
|
|
| def auto_detect_entities(text: str, src_lang: str) -> dict: |
| """Dispatcher: chọn NER engine theo ngôn ngữ nguồn.""" |
| if src_lang == "vi": |
| return auto_detect_entities_vi(text) |
| return auto_detect_entities_en(text) |
|
|
|
|
| |
| |
| |
| def mask_entities(text: str, entity_dict: dict) -> tuple: |
| """ |
| Thay tên riêng bằng placeholder trước khi đưa vào model dịch. |
| Trả về (text_đã_mask, {placeholder: tên_đích}) |
| |
| Ví dụ: |
| Input: "Lê Hoàng Sơn là kỹ sư" |
| Output: ("__E0__ là kỹ sư", {"__E0__": "Le Hoang Son"}) |
| """ |
| placeholder_map = {} |
| masked = text |
| sorted_names = sorted(entity_dict.keys(), key=len, reverse=True) |
|
|
| for i, name in enumerate(sorted_names): |
| |
| pattern = re.compile(re.escape(name), re.IGNORECASE | re.UNICODE) |
| if pattern.search(masked): |
| token = f"__E{i}__" |
| placeholder_map[token] = entity_dict[name] |
| masked = pattern.sub(token, masked) |
|
|
| return masked, placeholder_map |
|
|
|
|
| def restore_entities(text: str, placeholder_map: dict) -> str: |
| """Thay placeholder ngược lại về tên đúng sau khi dịch.""" |
| for token, correct_name in placeholder_map.items(): |
| text = text.replace(token, correct_name) |
| return text |
|
|
|
|
| |
| |
| |
| def smart_replace(text: str, entity_dict: dict) -> str: |
| """ |
| Post-processing: replace entity còn sót trong output |
| (những tên model tự dịch dù không qua mask). |
| """ |
| if not entity_dict: |
| return text |
| for ent in sorted(entity_dict, key=len, reverse=True): |
| pattern = re.compile( |
| rf"(?<!\w){re.escape(ent)}(?!\w)", |
| re.IGNORECASE | re.UNICODE, |
| ) |
| text = pattern.sub(entity_dict[ent], text) |
| return text |
|
|
|
|
| |
| |
| |
| def translate_logic(text: str, mode: str, use_kg: bool) -> str: |
| if not text or not text.strip(): |
| return "" |
|
|
| direction = "en->vi" if mode == "Anh -> Việt" else "vi->en" |
| src_lang = "en" if direction == "en->vi" else "vi" |
| target_lang = "vi" if direction == "en->vi" else "en" |
|
|
| if not use_kg: |
| try: |
| return get_model(direction)(text)[0]["translation_text"] |
| except Exception as e: |
| return f"[Lỗi model]: {e}" |
|
|
| try: |
| |
| |
| db_entities = fetch_entities_from_supabase(src_lang) |
| db_match = {k: v for k, v in db_entities.items() |
| if re.search(re.escape(k), text, re.IGNORECASE)} |
|
|
| |
| |
| ner_entities = auto_detect_entities(text, src_lang) |
|
|
| |
| all_entities = {**ner_entities, **db_match} |
| logger.info( |
| f"Tổng entity sẽ mask: {len(all_entities)} " |
| f"(DB={len(db_match)}, NER={len(ner_entities)})" |
| ) |
|
|
| |
| |
| masked_text, placeholder_map = mask_entities(text, all_entities) |
| logger.info(f"Text sau mask: {masked_text}") |
|
|
| raw_output = get_model(direction)(masked_text)[0]["translation_text"] |
| final_output = restore_entities(raw_output, placeholder_map) |
|
|
| |
| tgt_entities = fetch_entities_from_supabase(target_lang) |
| tgt_match = {k: v for k, v in tgt_entities.items() if k in final_output} |
| final_output = smart_replace(final_output, tgt_match) |
|
|
| return final_output |
|
|
| except Exception as e: |
| logger.error(f"Lỗi pipeline: {e}", exc_info=True) |
| try: |
| return get_model(direction)(text)[0]["translation_text"] |
| except Exception as e2: |
| return f"[Lỗi model]: {e2}" |
|
|
|
|
| |
| |
| |
| with gr.Blocks(theme=gr.themes.Soft(), title="AI Translation + KG") as demo: |
| gr.Markdown("# 🚀 Hệ thống Dịch thuật AI tích hợp Knowledge Graph + NER") |
| gr.Markdown( |
| "Dịch Anh ↔ Việt với pipeline 4 lớp bảo vệ thực thể:\n" |
| "**KG DB** → **NER tự động** → **Transliteration** → **NMT Model**" |
| ) |
|
|
| with gr.Row(): |
| with gr.Column(): |
| input_text = gr.Textbox( |
| lines=7, |
| label="Văn bản gốc", |
| placeholder="Nhập câu cần dịch tại đây…", |
| ) |
| mode = gr.Radio( |
| choices=["Anh -> Việt", "Việt -> Anh"], |
| label="Chiều dịch", |
| value="Anh -> Việt", |
| ) |
| use_kg = gr.Checkbox( |
| label="Bật Knowledge Graph + NER + Transliteration", |
| value=True, |
| ) |
| btn = gr.Button("🔄 Dịch ngay", variant="primary") |
|
|
| with gr.Column(): |
| output_text = gr.Textbox( |
| lines=7, |
| label="Kết quả dịch", |
| interactive=False, |
| ) |
| gr.Markdown( |
| "**Pipeline khi bật KG:**\n\n" |
| "1. 🗄️ Tra Knowledge Graph DB\n" |
| "2. 🔍 NER tự động phát hiện tên mới\n" |
| "3. 🔤 Transliteration: bỏ dấu tên riêng\n" |
| "4. 🤖 Model dịch (tên đã được bảo vệ)\n\n" |
| "→ Không còn lỗi `Lê Hoàng Sơn` → `Li Wongshan` ✅" |
| ) |
|
|
| with gr.Accordion("⚙️ Quản lý cache (admin)", open=False): |
| refresh_btn = gr.Button("🗑️ Xóa entity cache (reload từ Supabase)") |
| cache_status = gr.Textbox(label="Trạng thái", interactive=False) |
|
|
| def do_refresh(): |
| invalidate_cache() |
| return " Cache đã xóa. Entity sẽ reload ở request tiếp theo." |
|
|
| refresh_btn.click(fn=do_refresh, outputs=cache_status) |
|
|
| btn.click( |
| fn=translate_logic, |
| inputs=[input_text, mode, use_kg], |
| outputs=output_text, |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
