convert_to_answer.py

import argparse
import copy
import json
import os
import re
from collections import Counter
from glob import glob
from typing import Any, Dict, List, Optional, Tuple


DEFAULT_TEMPLATE_PATH = "outputs/predictions/track_1_test.json"
DEFAULT_PRED_PATH = "outputs/predictions/predict_ckpt2660/generated_predictions.jsonl"
DEFAULT_OUTPUT_PATH = "outputs/submissions/answers/track_1_test.json"


def load_json(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_jsonl(path: str) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    with open(path, "r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            raw = line.strip()
            if not raw:
                continue
            try:
                rows.append(json.loads(raw))
            except Exception:
                print(f"[WARN] skip invalid jsonl line: {line_no}")
    return rows


def normalize_level(level: Any) -> Optional[str]:
    if level is None:
        return None
    s = str(level).strip()
    if not s:
        return None

    m = {
        "poor": "Poor",
        "medium": "Medium",
        "good": "Good",
        "a": "A",
        "b": "B",
        "c": "C",
    }
    return m.get(s.lower(), s)


def score_to_submission_level(score: float, low_threshold: float = 5.0, high_threshold: float = 7.0) -> str:
    # 比赛映射：A=Poor(0~5), B=Medium(5~7), C=Good(7~10)
    if score < low_threshold:
        return "A"
    if score < high_threshold:
        return "B"
    return "C"


def to_submission_level(level: Any) -> Optional[str]:
    """

    Submission mapping required by Track-1 template:

    A -> Poor, B -> Medium, C -> Good

    """
    norm = normalize_level(level)
    if norm is None:
        return None
    m = {
        "Poor": "A",
        "Medium": "B",
        "Good": "C",
        "A": "A",
        "B": "B",
        "C": "C",
    }
    return m.get(norm)


def normalize_prediction_text(pred_row: Dict[str, Any]) -> str:
    # 兼容不同推理后端字段名。
    for k in ("predict", "response", "output", "text", "generation"):
        if k in pred_row:
            return str(pred_row.get(k, ""))
    return str(pred_row)


def extract_level_from_text(pred_text: str, crit_key: str) -> Optional[str]:
    pattern = rf'"{re.escape(crit_key)}"\s*:\s*\{{.*?"level"\s*:\s*"([^"]+)"'

    m = re.search(pattern, pred_text, flags=re.IGNORECASE | re.DOTALL)

    if not m:

        return None

    return to_submission_level(m.group(1))





def try_parse_predict_json(text: str) -> Optional[Dict[str, Any]]:

    text = str(text or "").strip()

    if not text:

        return None



    # 1) 直接解析

    try:

        obj = json.loads(text)

        if isinstance(obj, dict):

            return obj

    except Exception:

        pass



    # 2) 尝试提取从首个 { 到最后一个 } 的片段

    left = text.find("{")

    right = text.rfind("}")

    if left != -1 and right != -1 and left < right:

        snippet = text[left:right + 1]

        try:

            obj = json.loads(snippet)

            if isinstance(obj, dict):

                return obj

        except Exception:

            pass



    return None





def extract_answer(pred_obj: Optional[Dict[str, Any]], pred_text: str) -> Optional[str]:

    if pred_obj is not None:

        answer = str(pred_obj.get("answer", "")).strip().upper()

        if answer in {"A", "B", "C", "D"}:

            return answer



    m = re.search(r'"answer"\s*:\s*"([A-D])"', pred_text, re.IGNORECASE)

    if m:

        return m.group(1).upper()



    # 兼容 QA-only 推理：模型可能只输出单个字母（如 "C"）。

    raw = str(pred_text or "").strip().upper()

    if raw:

        # 情况1：整行仅有一个候选字母（允许尾随标点）。

        m = re.match(r"^\s*([A-D])(?:[\.\)\]:：]|\s)*$", raw)

        if m:

            return m.group(1)

        # 情况2：短文本中只出现唯一一个 A/B/C/D，且不包含常见 JSON/选项结构。

        if len(raw) <= 12 and ("{" not in raw) and ("\"" not in raw):

            hits = re.findall(r"[A-D]", raw)

            if len(hits) == 1:

                return hits[0]

    return None





def extract_total_score(pred_obj: Optional[Dict[str, Any]], pred_text: str) -> Optional[float]:

    val: Optional[float] = None

    if pred_obj is not None and "total_score" in pred_obj:

        try:

            val = float(pred_obj["total_score"])

        except Exception:

            val = None



    if val is None:

        m = re.search(r'"total_score"\s*:\s*([0-9]+(?:\.[0-9]+)?)', pred_text)

        if m:

            val = float(m.group(1))



    if val is None:

        return None



    return max(0.0, min(100.0, float(val)))





def merge_total_scores(total_votes: List[float], mode: str) -> Optional[int]:

    if not total_votes:

        return None



    vals = [float(v) for v in total_votes]

    if mode == "trim_mean" and len(vals) >= 3:

        vals = sorted(vals)[1:-1]

    merged = sum(vals) / len(vals)

    return max(0, min(100, int(round(merged))))





def parse_score_from_text(pred_text: str, crit_key: str) -> Optional[float]:

    pattern = rf'"{re.escape(crit_key)}"\s*:\s*\{{.*?"score"\s*:\s*([0-9]+(?:\.[0-9]+)?)'

    m = re.search(pattern, pred_text, flags=re.IGNORECASE | re.DOTALL)

    if not m:

        return None

    try:

        return float(m.group(1))

    except Exception:

        return None





def extract_one_criteria_level_and_score(

    crit_key: str,

    pred_obj: Optional[Dict[str, Any]],

    pred_text: str,

) -> Tuple[Optional[str], Optional[float]]:

    score: Optional[float] = None

    level: Optional[str] = None



    if isinstance(pred_obj, dict):

        src = pred_obj.get("criteria", {})

        if isinstance(src, dict) and crit_key in src:

            value = src.get(crit_key)

            if isinstance(value, dict):

                if "score" in value:

                    try:

                        score = float(value["score"])

                    except Exception:

                        score = None

                level = to_submission_level(value.get("level"))

            else:

                level = to_submission_level(value)



    if score is None:

        score = parse_score_from_text(pred_text, crit_key)



    if level is None:

        level = extract_level_from_text(pred_text, crit_key)



    if level not in {"A", "B", "C"}:

        level = None



    return level, score





def choose_majority_level(

    levels: List[str],

    scores: List[float],

    fallback_level: Optional[str],

    low_threshold: float,

    high_threshold: float,

) -> str:

    if levels:

        cnt = Counter(levels)

        top_n = max(cnt.values())

        top_levels = sorted([k for k, v in cnt.items() if v == top_n])

        if len(top_levels) == 1:

            return top_levels[0]



    # 平票时，用多次预测的均值 score 判定等级。

    if scores:

        mean_score = sum(scores) / len(scores)

        return score_to_submission_level(mean_score, low_threshold=low_threshold, high_threshold=high_threshold)



    if fallback_level in {"A", "B", "C"}:

        return fallback_level

    return "B"





def parse_weights(raw_weights: Optional[List[float]], n_models: int) -> List[float]:

    # 默认等权；若传入权重则要求和预测文件数一致。

    if raw_weights is None:

        return [1.0] * n_models

    if len(raw_weights) != n_models:

        raise ValueError(

            f"--weights length ({len(raw_weights)}) must equal number of prediction files ({n_models})."

        )

    for w in raw_weights:

        if w < 0:

            raise ValueError("weights must be non-negative.")

    # 全零没有意义，回退等权。

    if sum(raw_weights) == 0:

        return [1.0] * n_models

    return raw_weights





def load_thresholds(

    thresholds_json: str,

) -> Tuple[Dict[str, Dict[str, float]], Dict[str, float]]:

    """

    Accepts JSON in either format:

    1) {"criteria": {"Color Harmony": {"low": 4.9, "high": 7.1}}, "default": {"low":5,"high":7}}

    2) {"Color Harmony": {"low": 4.9, "high": 7.1}, ...}

    """

    default = {"low": 5.0, "high": 7.0}

    per_criteria: Dict[str, Dict[str, float]] = {}



    if not thresholds_json:

        return per_criteria, default

    if not os.path.exists(thresholds_json):

        print(f"[WARN] thresholds file not found: {thresholds_json}, fallback to default 5/7")

        return per_criteria, default



    obj = load_json(thresholds_json)

    if not isinstance(obj, dict):

        print(f"[WARN] invalid thresholds json format: {thresholds_json}, fallback to default 5/7")

        return per_criteria, default



    if "default" in obj and isinstance(obj.get("default"), dict):

        d = obj["default"]

        low = d.get("low", 5.0)

        high = d.get("high", 7.0)

        try:

            low_f = float(low)

            high_f = float(high)

            if low_f < high_f:

                default = {"low": low_f, "high": high_f}

        except Exception:

            pass



    src = obj.get("criteria") if isinstance(obj.get("criteria"), dict) else obj

    if isinstance(src, dict):

        for k, v in src.items():

            if not isinstance(v, dict):

                continue

            if "low" not in v or "high" not in v:

                continue

            try:

                low = float(v["low"])

                high = float(v["high"])

            except Exception:

                continue

            if low < high:

                per_criteria[str(k)] = {"low": low, "high": high}



    return per_criteria, default





def resolve_best_index(best_index: int, weights: List[float]) -> int:

    # best_index=-1 表示自动选择权重最高的模型作为平票时的优先模型。

    if best_index >= 0:

        if best_index >= len(weights):

            raise ValueError(f"--best_index out of range: {best_index}, num_models={len(weights)}")

        return best_index

    return max(range(len(weights)), key=lambda i: weights[i])





def extract_criteria_voting(

    template_item: Dict[str, Any],

    pred_objs: List[Optional[Dict[str, Any]]],

    pred_texts: List[str],

    per_criteria_thresholds: Dict[str, Dict[str, float]],

    default_thresholds: Dict[str, float],

) -> Dict[str, Dict[str, str]]:

    out: Dict[str, Dict[str, str]] = {}



    for crit_key, crit_val in template_item.get("criteria", {}).items():

        thresholds = per_criteria_thresholds.get(crit_key, default_thresholds)

        low = float(thresholds.get("low", 5.0))

        high = float(thresholds.get("high", 7.0))

        if not (low < high):

            low, high = 5.0, 7.0



        level_votes: List[str] = []

        score_votes: List[float] = []

        for obj, text in zip(pred_objs, pred_texts):

            level, score = extract_one_criteria_level_and_score(crit_key, obj, text)

            if level is not None:

                level_votes.append(level)

            if score is not None:

                score_votes.append(score)



        prev = str(crit_val.get("level", "")).strip() if isinstance(crit_val, dict) else ""

        final_level = choose_majority_level(

            level_votes,

            score_votes,

            prev,

            low_threshold=low,

            high_threshold=high,

        )

        out[crit_key] = {"level": final_level}



    return out





def pick_default_predictions_path() -> str:

    candidates = [

        "outputs/predictions/predict_ckpt2660/generated_predictions.jsonl",

        "outputs/predictions/generated_predictions.jsonl",

    ]

    # 自动兜底：在预测目录里找最近一次 generated_predictions.jsonl。

    dynamic = sorted(

        glob("outputs/predictions/**/generated_predictions.jsonl", recursive=True),

        key=lambda x: os.path.getmtime(x),

        reverse=True,

    )

    candidates = dynamic + candidates



    for p in candidates:

        if os.path.exists(p):

            return p

    return DEFAULT_PRED_PATH





def choose_weighted_answer(

    votes_by_model: List[Optional[str]],

    weights: List[float],

    best_index: int,

    tie_break_answer: Optional[str],

) -> str:

    """

    更合理的答案融合策略：

    1) 加权投票（按各变体可靠性权重）

    2) 若平票且提供 tie-break 结果，则优先用 tie-break

    3) 若仍平票，采用最佳变体(best_index)在平票选项中的答案

    4) 最后才做稳定兜底（字母序）

    """

    label_scores = {"A": 0.0, "B": 0.0, "C": 0.0, "D": 0.0}

    for i, ans in enumerate(votes_by_model):

        if ans in label_scores:

            label_scores[ans] += weights[i]



    max_score = max(label_scores.values())

    if max_score <= 0:

        return "A"



    tied = sorted([k for k, v in label_scores.items() if v == max_score])

    if len(tied) == 1:

        return tied[0]



    if tie_break_answer in tied:

        return tie_break_answer  # 用专门 tie-break 结果判平票



    best_vote = votes_by_model[best_index]

    if best_vote in tied:

        return str(best_vote)



    return tied[0]





def main() -> None:

    parser = argparse.ArgumentParser()

    parser.add_argument("--template_json", type=str, default=DEFAULT_TEMPLATE_PATH)

    parser.add_argument(

        "--predictions_jsonl",

        type=str,

        nargs="+",

        default=None,

        help="One or more generated_predictions.jsonl paths for voting.",

    )

    parser.add_argument(

        "--weights",

        type=float,

        nargs="+",

        default=None,

        help="Optional weights for prediction files (same length as --predictions_jsonl).",

    )

    parser.add_argument(

        "--best_index",

        type=int,

        default=-1,

        help="Best model index for tie fallback. -1 means auto argmax(weights).",

    )

    parser.add_argument(

        "--tie_break_jsonl",

        type=str,

        default="",

        help="Optional tie-break predictions file used only when weighted vote ties.",

    )

    parser.add_argument(

        "--thresholds_json",

        type=str,

        default="",

        help="Optional per-criterion thresholds json for score->A/B/C mapping.",

    )

    parser.add_argument(

        "--total_score_fusion",

        type=str,

        default="mean",

        choices=["mean", "trim_mean"],

        help="Fusion mode for total_score across multi-prompt predictions.",

    )

    parser.add_argument("--output_json", type=str, default=DEFAULT_OUTPUT_PATH)

    args = parser.parse_args()



    pred_paths = args.predictions_jsonl if args.predictions_jsonl else [pick_default_predictions_path()]

    pred_paths = [p for p in pred_paths if str(p).strip()]

    if not pred_paths:

        raise ValueError("No predictions_jsonl provided or discovered.")



    template_data = load_json(args.template_json)

    pred_sets = [load_jsonl(p) for p in pred_paths]

    weights = parse_weights(args.weights, len(pred_sets))

    best_index = resolve_best_index(args.best_index, weights)

    per_criteria_thresholds, default_thresholds = load_thresholds(args.thresholds_json)

    tie_break_rows: List[Dict[str, Any]] = []

    if args.tie_break_jsonl:

        tie_break_rows = load_jsonl(args.tie_break_jsonl)



    if not isinstance(template_data, list):

        raise ValueError("template_json must be a list.")



    print(f"[INFO] template items: {len(template_data)}")

    for p, rows in zip(pred_paths, pred_sets):

        print(f"[INFO] prediction rows: {len(rows)} ({p})")

    print(f"[INFO] answer weights: {weights}")

    print(f"[INFO] answer best_index: {best_index}")

    print(f"[INFO] criteria thresholds default: low={default_thresholds['low']}, high={default_thresholds['high']}")

    if per_criteria_thresholds:

        print(f"[INFO] criteria thresholds loaded: {len(per_criteria_thresholds)}")

    if args.tie_break_jsonl:

        print(f"[INFO] tie_break rows: {len(tie_break_rows)} ({args.tie_break_jsonl})")



    final_submission: List[Dict[str, Any]] = []

    parsed_ok = 0

    filled = 0



    for i, item in enumerate(template_data):

        new_item = copy.deepcopy(item)

        # 收集每次预测在第 i 条样本上的结果。

        row_texts: List[str] = []

        row_objs: List[Optional[Dict[str, Any]]] = []

        answer_votes_by_model: List[Optional[str]] = []

        total_votes: List[int] = []



        for rows in pred_sets:

            if i >= len(rows):

                answer_votes_by_model.append(None)

                continue

            pred_text = normalize_prediction_text(rows[i])

            pred_obj = try_parse_predict_json(pred_text)

            if pred_obj is not None:

                parsed_ok += 1

            row_texts.append(pred_text)

            row_objs.append(pred_obj)



            ans = extract_answer(pred_obj, pred_text)

            answer_votes_by_model.append(ans)



            ts = extract_total_score(pred_obj, pred_text)

            if ts is not None:

                total_votes.append(ts)



        if not row_texts:

            final_submission.append(new_item)

            continue



        new_item["criteria"] = extract_criteria_voting(

            new_item,

            row_objs,

            row_texts,

            per_criteria_thresholds=per_criteria_thresholds,

            default_thresholds=default_thresholds,

        )



        # total_score 用多次预测均值，减少单次抖动。

        merged_total = merge_total_scores(total_votes, args.total_score_fusion)

        if merged_total is not None:

            new_item["total_score"] = max(0, min(100, merged_total))



        tie_break_answer: Optional[str] = None

        if i < len(tie_break_rows):

            tb_text = normalize_prediction_text(tie_break_rows[i])

            tb_obj = try_parse_predict_json(tb_text)

            tie_break_answer = extract_answer(tb_obj, tb_text)



        new_item["answer"] = choose_weighted_answer(

            votes_by_model=answer_votes_by_model,

            weights=weights,

            best_index=best_index,

            tie_break_answer=tie_break_answer,

        )



        final_submission.append(new_item)

        filled += 1



    os.makedirs(os.path.dirname(args.output_json), exist_ok=True)

    with open(args.output_json, "w", encoding="utf-8") as f:

        json.dump(final_submission, f, ensure_ascii=False, indent=2)



    parsed_total = filled * len(pred_paths)

    print(f"[INFO] parsed predict json ok: {parsed_ok}/{parsed_total}")

    print(f"[INFO] saved submission: {args.output_json}")





if __name__ == "__main__":

    main()