inference.py

import json
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd


def load_reranker(bundle_dir):
    bundle_dir = Path(bundle_dir)

    model_path = bundle_dir / "sql_reranker_lightgbm.txt"
    feature_cols_path = bundle_dir / "feature_columns.json"

    if not model_path.exists():
        raise FileNotFoundError(f"Missing model file: {model_path}")

    if not feature_cols_path.exists():
        raise FileNotFoundError(f"Missing feature columns file: {feature_cols_path}")

    model = lgb.Booster(model_file=str(model_path))

    with open(feature_cols_path, "r", encoding="utf-8") as f:
        feature_cols = json.load(f)

    return model, feature_cols


def select_best_sql_from_feature_rows(candidate_feature_rows, model, feature_cols):
    if not candidate_feature_rows:
        return {
            "best_sql": "",
            "best_score": None,
            "best_index": None,
            "all_scores": [],
        }

    df = pd.DataFrame(candidate_feature_rows)

    if "candidate_sql" not in df.columns:
        raise ValueError("candidate_feature_rows must include candidate_sql.")

    X = df.reindex(columns=feature_cols, fill_value=0)
    X = X.apply(pd.to_numeric, errors="coerce").fillna(0)

    scores = model.predict(X)
    best_idx = int(np.argmax(scores))

    return {
        "best_sql": df.iloc[best_idx]["candidate_sql"],
        "best_score": float(scores[best_idx]),
        "best_index": best_idx,
        "all_scores": scores.tolist(),
    }