text2sql-sqam.py

import evaluate
import datasets

from text2sql_eval.metrics.sqam_wrapper import sqam_score


_DESCRIPTION = "SQAM (Structural Query Alignment Metric) for SQL. Returns mean score in [0, 1]."


def _to_str(x):
    # Evaluate sometimes passes references as list-of-lists; we accept both.
    if isinstance(x, (list, tuple)):
        return x[0] if x else ""
    return "" if x is None else str(x)


class SQAM(evaluate.Metric):
    def _info(self):
        return evaluate.MetricInfo(
            description=_DESCRIPTION,
            citation="SQAM: https://github.com/ezzini/SQAM",
            features=datasets.Features(
                {
                    "predictions": datasets.Value("string"),
                    "references": datasets.Value("string"),
                }
            ),
        )

    def _compute(self, predictions, references):
        scores = []
        for p, r in zip(predictions, references):
            scores.append(float(sqam_score(_to_str(p), _to_str(r))))
        mean = sum(scores) / len(scores) if scores else 0.0
        return {"sqam": mean}