labor-arbitration-system/backend/tools/evaluate_extractor.py

from __future__ import annotations

import json
from collections import defaultdict
from dataclasses import dataclass
from typing import Any

from app.extractor import RuleBasedLaborExtractor


def _norm(v: Any) -> Any:
    if v is None:
        return None
    if isinstance(v, str):
        return v.strip()
    return v


def _as_set(v: Any) -> set[str]:
    if v is None:
        return set()
    if isinstance(v, list):
        return set([str(x).strip() for x in v if str(x).strip()])
    return {str(v).strip()}


@dataclass
class FieldPRF:
    tp: int = 0
    fp: int = 0
    fn: int = 0

    def add(self, gold: Any, pred: Any) -> None:
        gset = _as_set(gold)
        pset = _as_set(pred)
        self.tp += len(gset & pset)
        self.fp += len(pset - gset)
        self.fn += len(gset - pset)

    def metrics(self) -> dict[str, float]:
        p = self.tp / (self.tp + self.fp) if (self.tp + self.fp) else 0.0
        r = self.tp / (self.tp + self.fn) if (self.tp + self.fn) else 0.0
        f1 = 2 * p * r / (p + r) if (p + r) else 0.0
        return {"precision": round(p, 4), "recall": round(r, 4), "f1": round(f1, 4)}


def run_eval(samples: list[dict[str, Any]]) -> dict[str, Any]:
    extractor = RuleBasedLaborExtractor()
    stats: dict[str, FieldPRF] = defaultdict(FieldPRF)

    for s in samples:
        pred = extractor.extract(s["text"])
        gold = s["gold"]
        for field in gold.keys():
            if field not in pred:
                continue
            if field == "claims":
                # claims 是结构体：只评 items 和 amount_total
                stats["claims.items"].add(gold["claims"].get("items"), pred["claims"].get("items"))
                stats["claims.amount_total"].add(gold["claims"].get("amount_total"), pred["claims"].get("amount_total"))
            else:
                stats[field].add(_norm(gold.get(field)), _norm(pred.get(field)))

    report = {k: v.metrics() | {"tp": v.tp, "fp": v.fp, "fn": v.fn} for k, v in sorted(stats.items())}
    # micro 平均（合计）
    total = FieldPRF()
    for v in stats.values():
        total.tp += v.tp
        total.fp += v.fp
        total.fn += v.fn
    report["_micro_avg"] = total.metrics() | {"tp": total.tp, "fp": total.fp, "fn": total.fn}
    return report


def demo_samples() -> list[dict[str, Any]]:
    """
    至少 10 条模拟样例（可替换为你的真实标注集）。
    """
    return [
        {
            "text": "申请人：张三\n被申请人：北京某某有限公司\n单位性质：企业\n岗位：销售\n入职时间：2020年1月1日\n离职时间：2022-03-01\n月工资：8000元\n仲裁请求：1. 支付拖欠工资8000元；2. 支付加班费2000元。\n依据《劳动合同法》第三十条。",
            "gold": {
                "applicant_name": "张三",
                "respondent_name": "北京某某有限公司",
                "employer_nature": "企业",
                "worker_position": "销售",
                "case_cause": "工资报酬",
                "entry_date": "2020-01-01",
                "leave_date": "2022-03-01",
                "month_salary": 8000.0,
                "overtime_desc": "支付加班费2000元",
                "termination_reason": None,
                "claims": {"items": ["支付拖欠工资8000元", "支付加班费2000元"], "amount_total": 10000.0},
                "law_refs": ["《劳动合同法》第三十条", "《劳动合同法》"],
            },
        },
        {
            "text": "申请人：李四，被申请人：上海某医院。岗位：护士。于2020.02.02入职，于2021.12.31解除。请求经济补偿八千元。引用《劳动合同法》。",
            "gold": {
                "applicant_name": "李四",
                "respondent_name": "上海某医院",
                "employer_nature": "事业单位",
                "worker_position": "护士",
                "case_cause": "经济补偿",
                "entry_date": "2020-02-02",
                "leave_date": "2021-12-31",
                "month_salary": None,
                "overtime_desc": None,
                "termination_reason": None,
                "claims": {"items": [], "amount_total": None},
                "law_refs": ["《劳动合同法》"],
            },
        },
        {
            "text": "申请人：王五\n被申请人：某某科技股份有限公司\n入职日期2021/6/1，离职日期2023/1/15。\n月薪为8,000元。\n因绩效不达标被辞退。\n请求：违法解除劳动合同赔偿金96000元。\n依据劳动合同法第八十七条。",
            "gold": {
                "applicant_name": "王五",
                "respondent_name": "某某科技股份有限公司",
                "employer_nature": "企业",
                "worker_position": None,
                "case_cause": "违法解除劳动合同",
                "entry_date": "2021-06-01",
                "leave_date": "2023-01-15",
                "month_salary": 8000.0,
                "overtime_desc": None,
                "termination_reason": "绩效不达标",
                "claims": {"items": ["违法解除劳动合同赔偿金96000元"], "amount_total": 96000.0},
                "law_refs": ["劳动合同法第八十七条"],
            },
        },
        {
            "text": "申请人：赵六\n被申请人：某个体工商户\n单位性质：个体工商户\n岗位：厨师\n入职时间：2019-5-10\n离职时间：2020-6-1\n主张加班费3000元，休息日加班无调休。\n法律依据：《劳动争议调解仲裁法》。",
            "gold": {
                "applicant_name": "赵六",
                "respondent_name": "某个体工商户",
                "employer_nature": "个人",
                "worker_position": "厨师",
                "case_cause": "加班费",
                "entry_date": "2019-05-10",
                "leave_date": "2020-06-01",
                "month_salary": None,
                "overtime_desc": "主张加班费3000元，休息日加班无调休",
                "termination_reason": None,
                "claims": {"items": [], "amount_total": None},
                "law_refs": ["《劳动争议调解仲裁法》"],
            },
        },
        {
            "text": "申请人：孙七\n被申请人：某某有限公司\n工伤待遇：一次性伤残补助金50000元。\n于2022年7月7日入职，2023年7月7日离职。\n依据《工伤保险条例》。",
            "gold": {
                "applicant_name": "孙七",
                "respondent_name": "某某有限公司",
                "employer_nature": "企业",
                "worker_position": None,
                "case_cause": "工伤待遇",
                "entry_date": "2022-07-07",
                "leave_date": "2023-07-07",
                "month_salary": None,
                "overtime_desc": None,
                "termination_reason": None,
                "claims": {"items": [], "amount_total": None},
                "law_refs": ["《工伤保险条例》"],
            },
        },
        # 其余 5 条：覆盖不同日期/金额/案由写法
        {
            "text": "申请人：周八 被申请人：某某集团有限公司 岗位：工程师 入职：2020年12月5日 离职：2022年2月1日 月薪12000元 请求加班费10000元。",
            "gold": {
                "applicant_name": "周八",
                "respondent_name": "某某集团有限公司",
                "employer_nature": "企业",
                "worker_position": "工程师",
                "case_cause": "加班费",
                "entry_date": "2020-12-05",
                "leave_date": "2022-02-01",
                "month_salary": 12000.0,
                "overtime_desc": "请求加班费10000元",
                "termination_reason": None,
                "claims": {"items": [], "amount_total": None},
                "law_refs": [],
            },
        },
        {
            "text": "申请人：钱九\n被申请人：某某学校\n单位性质：事业单位\n请求：支付拖欠工资二万元。\n入职时间2021.01.01，解除时间2021.06.30。",
            "gold": {
                "applicant_name": "钱九",
                "respondent_name": "某某学校",
                "employer_nature": "事业单位",
                "worker_position": None,
                "case_cause": "工资报酬",
                "entry_date": "2021-01-01",
                "leave_date": "2021-06-30",
                "month_salary": None,
                "overtime_desc": None,
                "termination_reason": None,
                "claims": {"items": ["支付拖欠工资二万元"], "amount_total": 20000.0},
                "law_refs": [],
            },
        },
        {
            "text": "申请人：吴十\n被申请人：某某有限责任公司\n解除原因：未签订劳动合同。\n请求经济补偿12000元。\n法律条款：《劳动合同法》第四十六条。",
            "gold": {
                "applicant_name": "吴十",
                "respondent_name": "某某有限责任公司",
                "employer_nature": "企业",
                "worker_position": None,
                "case_cause": "经济补偿",
                "entry_date": None,
                "leave_date": None,
                "month_salary": None,
                "overtime_desc": None,
                "termination_reason": "未签订劳动合同",
                "claims": {"items": ["经济补偿12000元"], "amount_total": 12000.0},
                "law_refs": ["《劳动合同法》第四十六条", "《劳动合同法》"],
            },
        },
        {
            "text": "申请人：郑十一\n被申请人：某某公司\n岗位：司机\n加班事实：长期每天加班至22点，未支付加班工资。\n请求加班费三千元。\n",
            "gold": {
                "applicant_name": "郑十一",
                "respondent_name": "某某公司",
                "employer_nature": "企业",
                "worker_position": "司机",
                "case_cause": "加班费",
                "entry_date": None,
                "leave_date": None,
                "month_salary": None,
                "overtime_desc": "加班事实：长期每天加班至22点，未支付加班工资",
                "termination_reason": None,
                "claims": {"items": ["加班费三千元"], "amount_total": 3000.0},
                "law_refs": [],
            },
        },
        {
            "text": "申请人：冯十二\n被申请人：某某厂\n请求：支付工资差额5000元；支付经济补偿5000元。\n依据《劳动合同法》。",
            "gold": {
                "applicant_name": "冯十二",
                "respondent_name": "某某厂",
                "employer_nature": "企业",
                "worker_position": None,
                "case_cause": "工资报酬",
                "entry_date": None,
                "leave_date": None,
                "month_salary": None,
                "overtime_desc": None,
                "termination_reason": None,
                "claims": {"items": ["支付工资差额5000元", "支付经济补偿5000元"], "amount_total": 10000.0},
                "law_refs": ["《劳动合同法》"],
            },
        },
    ]


if __name__ == "__main__":
    # 可选：用 --dataset 读取你的真实标注集（JSON）
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", default="", help="可选：标注集 JSON 文件路径")
    args = parser.parse_args()

    if args.dataset:
        samples = json.loads(open(args.dataset, "r", encoding="utf-8").read())
    else:
        samples = demo_samples()

    report = run_eval(samples)
    print(json.dumps(report, ensure_ascii=False, indent=2))