inference.py

"""
IC Design Content Firewall (ZH) — Inference Module

Loads the LoRA adapter from Hugging Face Hub on top of Qwen2.5-1.5B-Instruct,
produces structured JSON for IC-design content moderation.

Output schema:
  {
    "content_categories": [...],          # multi-label, 9 types
    "risk_level": "L0/L1/L2/L3",
    "sensitive_entities": [{"type", "value", "reason"}, ...],
    "reasoning": {"category_reason", "risk_reason"}
  }

Two universal filters (Inference safeguards):
  Filter 1 - Schema validation : 確保 enum 合法
  Filter 2 - Provenance check  : 確保 entity value 必在原文（防幻覺）
"""

import json
import re
from typing import Any, Dict, Optional

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer


# =============================================================================
# 預設 model 路徑（HF Hub）
# =============================================================================

DEFAULT_BASE_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
DEFAULT_ADAPTER    = "GOSHUNCLE/ic-firewall-zh"

MAX_NEW_TOKENS = 512


# =============================================================================
# System prompt
# =============================================================================

SYSTEM_PROMPT = """你是 IC 設計業內容防火牆分析專家。請分析輸入文字，輸出嚴格 JSON。

任務：
1. content_categories：從以下 9 類中選擇相關的（multi-label）：
   RTL, CUSTOMER, QUOTE, VENDOR, PROCESS, SCHEDULE, TESTING, INTERNAL, PUBLIC
2. risk_level：選擇 L0/L1/L2/L3 之一
3. sensitive_entities：標註所有機敏實體（不取代，僅標註）
   每筆含 type、value（必須逐字在原文）、reason（10-20 字）
4. reasoning：含 category_reason 與 risk_reason（各 30-60 字）

實體類型：
CUSTOMER, PROJECT, VENDOR, PRICE, PROCESS_NODE,
MODULE_NAME, IP_BLOCK, YIELD_NUMBER, SPEC_PARAM

輸出格式：
{"content_categories":[...],"risk_level":"L?","sensitive_entities":[{"type":"...","value":"...","reason":"..."}],"reasoning":{"category_reason":"...","risk_reason":"..."}}

規則：
1. value 必須逐字出現在原文
2. 不要保險全選類別，只標 actually 相關的
3. 只回傳 JSON，無其他說明文字"""


# =============================================================================
# 合法 enum
# =============================================================================

VALID_CATEGORIES = {
    "RTL", "CUSTOMER", "QUOTE", "VENDOR", "PROCESS",
    "SCHEDULE", "TESTING", "INTERNAL", "PUBLIC"
}
VALID_RISKS = {"L0", "L1", "L2", "L3"}
VALID_ENTITY_TYPES = {
    "CUSTOMER", "PROJECT", "VENDOR", "PRICE", "PROCESS_NODE",
    "MODULE_NAME", "IP_BLOCK", "YIELD_NUMBER", "SPEC_PARAM"
}


# =============================================================================
# JSON 解析輔助
# =============================================================================

def extract_json(raw_text: str) -> Optional[Dict[str, Any]]:
    text = raw_text.strip()
    try:
        return json.loads(text)
    except Exception:
        pass
    s, e = text.find("{"), text.rfind("}")
    if s >= 0 and e > s:
        try:
            return json.loads(text[s:e + 1])
        except Exception:
            pass
    cleaned = re.sub(r"```(?:json)?", "", text).strip().replace("```", "").strip()
    try:
        return json.loads(cleaned)
    except Exception:
        return None


# =============================================================================
# Filter 1 — Schema validation
# =============================================================================

def filter_schema(parsed: Dict[str, Any]) -> Dict[str, Any]:
    cleaned = {
        "content_categories": [],
        "risk_level": "L1",
        "sensitive_entities": [],
        "reasoning": {"category_reason": "", "risk_reason": ""},
    }

    cats = parsed.get("content_categories", [])
    if isinstance(cats, list):
        seen = set()
        for c in cats:
            if isinstance(c, str) and c in VALID_CATEGORIES and c not in seen:
                cleaned["content_categories"].append(c)
                seen.add(c)
    if not cleaned["content_categories"]:
        cleaned["content_categories"] = ["INTERNAL"]

    risk = parsed.get("risk_level", "")
    if isinstance(risk, str) and risk in VALID_RISKS:
        cleaned["risk_level"] = risk

    entities = parsed.get("sensitive_entities", [])
    if isinstance(entities, list):
        for e in entities:
            if not isinstance(e, dict):
                continue
            t = e.get("type", "")
            v = e.get("value", "")
            r = e.get("reason", "")
            if not (isinstance(t, str) and t in VALID_ENTITY_TYPES):
                continue
            if not (isinstance(v, str) and v):
                continue
            cleaned["sensitive_entities"].append({
                "type": t, "value": v, "reason": str(r),
            })

    rsn = parsed.get("reasoning", {})
    if isinstance(rsn, dict):
        cr = rsn.get("category_reason", "")
        rr = rsn.get("risk_reason", "")
        if isinstance(cr, str):
            cleaned["reasoning"]["category_reason"] = cr
        if isinstance(rr, str):
            cleaned["reasoning"]["risk_reason"] = rr

    return cleaned


# =============================================================================
# Filter 2 — Provenance check
# =============================================================================

def _normalize(s: str) -> str:
    return s.replace(" ", "").replace("　", "").replace("\t", "")


def value_in_text(value: str, text: str) -> bool:
    if not value:
        return False
    if value in text:
        return True
    return _normalize(value) in _normalize(text)


def filter_provenance(parsed: Dict[str, Any], input_text: str) -> Dict[str, Any]:
    kept = [e for e in parsed.get("sensitive_entities", [])
            if value_in_text(e.get("value", ""), input_text)]
    parsed["sensitive_entities"] = kept
    return parsed


# =============================================================================
# Detector
# =============================================================================

class Detector:
    def __init__(self,
                 base_model_path: str = DEFAULT_BASE_MODEL,
                 adapter_path: Optional[str] = DEFAULT_ADAPTER,
                 use_adapter: bool = True):
        print(f"[Detector] Loading tokenizer from {base_model_path}")
        self.tokenizer = AutoTokenizer.from_pretrained(base_model_path)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        print(f"[Detector] Loading base model (fp32 CPU)...")
        self.model = AutoModelForCausalLM.from_pretrained(
            base_model_path,
            torch_dtype=torch.float32,
            device_map="cpu",
        )

        if use_adapter and adapter_path:
            print(f"[Detector] Applying LoRA adapter from {adapter_path}")
            self.model = PeftModel.from_pretrained(self.model, adapter_path)
        self.model.eval()
        print(f"[Detector] Ready.")

    def generate_raw(self, text: str) -> str:
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user",   "content": text},
        ]
        prompt_text = self.tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True,
        )
        inputs = self.tokenizer(prompt_text, return_tensors="pt", add_special_tokens=False)
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
                num_beams=1,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
            )
        new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        return self.tokenizer.decode(new_tokens, skip_special_tokens=True)

    @staticmethod
    def post_process(raw_output: str, input_text: str,
                      apply_filters: bool = True) -> Dict[str, Any]:
        parsed = extract_json(raw_output)
        if parsed is None:
            return {
                "content_categories": ["INTERNAL"] if apply_filters else [],
                "risk_level": "L1" if apply_filters else "",
                "sensitive_entities": [],
                "reasoning": {
                    "category_reason": "(JSON parse failed)",
                    "risk_reason": "(JSON parse failed)",
                },
                "_parse_failed": True,
            }
        if apply_filters:
            result = filter_schema(parsed)
            result = filter_provenance(result, input_text)
        else:
            # Lenient mode
            result = {
                "content_categories": parsed.get("content_categories", []) if isinstance(parsed.get("content_categories"), list) else [],
                "risk_level": parsed.get("risk_level", "") if isinstance(parsed.get("risk_level"), str) else "",
                "sensitive_entities": parsed.get("sensitive_entities", []) if isinstance(parsed.get("sensitive_entities"), list) else [],
                "reasoning": parsed.get("reasoning", {"category_reason": "", "risk_reason": ""}) if isinstance(parsed.get("reasoning"), dict) else {"category_reason": "", "risk_reason": ""},
            }
        return result

    def detect(self, text: str, apply_filters: bool = True) -> Dict[str, Any]:
        raw = self.generate_raw(text)
        return self.post_process(raw, text, apply_filters=apply_filters)