evaluation_logic.py

import os
import sys
from pathlib import Path
from datetime import datetime
import json
import traceback
import uuid
from huggingface_hub import CommitScheduler

current_dir = Path(__file__).resolve().parent
duckdb_nsql_dir = current_dir / 'duckdb-nsql'
eval_dir = duckdb_nsql_dir / 'eval'
sys.path.extend([str(current_dir), str(duckdb_nsql_dir), str(eval_dir)])

from eval.predict import get_manifest, DefaultLoader, PROMPT_FORMATTERS, generate_sql
from eval.evaluate import evaluate, compute_metrics, get_to_print
from eval.evaluate import test_suite_evaluation, read_tables_json
from eval.schema import TextToSQLParams, Table

AVAILABLE_PROMPT_FORMATS = list(PROMPT_FORMATTERS.keys())

prediction_folder = Path("prediction_results/")
evaluation_folder = Path("evaluation_results/")

file_uuid = uuid.uuid4()

prediction_scheduler = CommitScheduler(
    repo_id="sql-console/duckdb-nsql-predictions",
    repo_type="dataset",
    folder_path=prediction_folder,
    path_in_repo="data",
    every=10,
)

evaluation_scheduler = CommitScheduler(
    repo_id="sql-console/duckdb-nsql-scores",
    repo_type="dataset",
    folder_path=evaluation_folder,
    path_in_repo="data",
    every=10,
)

def save_prediction(inference_api, model_name, prompt_format, question, generated_sql):
    prediction_file = prediction_folder / f"prediction_{file_uuid}.json"
    prediction_folder.mkdir(parents=True, exist_ok=True)
    with prediction_scheduler.lock:
        with prediction_file.open("a") as f:
            json.dump({
                "inference_api": inference_api,
                "model_name": model_name,
                "prompt_format": prompt_format,
                "question": question,
                "generated_sql": generated_sql,
                "timestamp": datetime.now().isoformat()
            }, f)

def save_evaluation(inference_api, model_name, prompt_format, custom_prompt, metrics):
    evaluation_file = evaluation_folder / f"evaluation_{file_uuid}.json"
    evaluation_folder.mkdir(parents=True, exist_ok=True)

    # Extract and flatten the category-specific execution metrics
    categories = ['easy', 'medium', 'hard', 'duckdb', 'ddl', 'all']
    flattened_metrics = {
        "inference_api": inference_api,
        "model_name": model_name,
        "prompt_format": prompt_format,
        "custom_prompt": str(custom_prompt) if prompt_format.startswith("custom") else "",
        "timestamp": datetime.now().isoformat()
    }

    # Flatten each category's metrics into separate columns
    for category in categories:
        if category in metrics['exec']:
            category_metrics = metrics['exec'][category]
            flattened_metrics[f"{category}_count"] = category_metrics['count']
            flattened_metrics[f"{category}_execution_accuracy"] = category_metrics['exec']
        else:
            flattened_metrics[f"{category}_count"] = 0
            flattened_metrics[f"{category}_execution_accuracy"] = 0.0

    with evaluation_scheduler.lock:
        with evaluation_file.open("a") as f:
            json.dump(flattened_metrics, f)
            f.write('\n')

def run_prediction(inference_api, model_name, prompt_format, custom_prompt, output_file):
    dataset_path = str(eval_dir / "data/dev.json")
    table_meta_path = str(eval_dir / "data/tables.json")
    stop_tokens = ['`<|dummy|>`']
    max_tokens = 1000
    temperature = 0
    num_beams = -1
    manifest_client = inference_api
    manifest_engine = model_name
    manifest_connection = "http://localhost:5000"
    overwrite_manifest = True
    parallel = False

    yield "Starting prediction..."

    try:
        # Initialize necessary components
        data_formatter = DefaultLoader()
        if prompt_format.startswith("custom"):
            prompt_formatter_cls = PROMPT_FORMATTERS["custom"]
            prompt_formatter_cls.PROMPT_TEMPLATE = custom_prompt
            prompt_formatter = prompt_formatter_cls()
        else:
            prompt_formatter = PROMPT_FORMATTERS[prompt_format]()

        # Load manifest
        manifest = get_manifest(
            manifest_client=manifest_client,
            manifest_connection=manifest_connection,
            manifest_engine=manifest_engine,
        )

        # Load data
        data = data_formatter.load_data(dataset_path)
        db_to_tables = data_formatter.load_table_metadata(table_meta_path)

        # Prepare input for generate_sql
        text_to_sql_inputs = []
        for input_question in data:
            question = input_question["question"]
            db_id = input_question.get("db_id", "none")
            if db_id != "none":
                table_params = list(db_to_tables.get(db_id, {}).values())
            else:
                table_params = []

            text_to_sql_inputs.append(TextToSQLParams(
                instruction=question,
                database=db_id,
                tables=table_params,
            ))

        # Generate SQL
        generated_sqls = generate_sql(
            manifest=manifest,
            text_to_sql_in=text_to_sql_inputs,
            retrieved_docs=[[] for _ in text_to_sql_inputs],
            prompt_formatter=prompt_formatter,
            stop_tokens=stop_tokens,
            overwrite_manifest=overwrite_manifest,
            max_tokens=max_tokens,
            temperature=temperature,
            num_beams=num_beams,
            parallel=parallel
        )

        # Save results
        output_file.parent.mkdir(parents=True, exist_ok=True)
        with output_file.open('w') as f:
            for original_data, (sql, _) in zip(data, generated_sqls):
                output = {**original_data, "pred": sql}
                json.dump(output, f)
                f.write('\n')

                # Save prediction to dataset
                save_prediction(inference_api, model_name, prompt_format, original_data["question"], sql)

        yield f"Prediction completed. Results saved to {output_file}"
    except Exception as e:
        yield f"Prediction failed with error: {str(e)}"
        yield f"Error traceback: {traceback.format_exc()}"

def run_evaluation(inference_api, model_name, prompt_format="duckdbinstgraniteshort", custom_prompt=None):
    if "OPENROUTER_API_KEY" not in os.environ:
        yield "Error: OPENROUTER_API_KEY not found in environment variables."
        return
    if "HF_TOKEN" not in os.environ:
        yield "Error: HF_TOKEN not found in environment variables."
        return

    try:
        # Set up the arguments
        dataset_path = str(eval_dir / "data/dev.json")
        table_meta_path = str(eval_dir / "data/tables.json")
        output_dir = eval_dir / "output"

        yield f"Using model: {model_name}"
        yield f"Using prompt format: {prompt_format}"

        if prompt_format == "custom":
            prompt_format = prompt_format+"_"+str(abs(hash(custom_prompt)) % (10 ** 8))

        output_file = output_dir / f"{prompt_format}_0docs_{model_name.replace('/', '_')}_dev_{datetime.now().strftime('%y-%m-%d')}.json"

        # Ensure the output directory exists
        output_dir.mkdir(parents=True, exist_ok=True)

        if output_file.exists():
            yield f"Prediction file already exists: {output_file}"
            yield "Skipping prediction step and proceeding to evaluation."
        else:
            # Run prediction
            for output in run_prediction(inference_api, model_name, prompt_format, custom_prompt, output_file):
                yield output

        # Run evaluation
        yield "Starting evaluation..."

        # Set up evaluation arguments
        gold_path = Path(dataset_path)
        db_dir = str(eval_dir / "data/databases/")
        tables_path = Path(table_meta_path)

        kmaps = test_suite_evaluation.build_foreign_key_map_from_json(str(tables_path))
        db_schemas = read_tables_json(str(tables_path))

        gold_sqls_dict = json.load(gold_path.open("r", encoding="utf-8"))
        pred_sqls_dict = [json.loads(l) for l in output_file.open("r").readlines()]

        gold_sqls = [p.get("query", p.get("sql", "")) for p in gold_sqls_dict]
        setup_sqls = [p["setup_sql"] for p in gold_sqls_dict]
        validate_sqls = [p["validation_sql"] for p in gold_sqls_dict]
        gold_dbs = [p.get("db_id", p.get("db", "")) for p in gold_sqls_dict]
        pred_sqls = [p["pred"] for p in pred_sqls_dict]
        categories = [p.get("category", "") for p in gold_sqls_dict]

        yield "Computing metrics..."
        metrics = compute_metrics(
            gold_sqls=gold_sqls,
            pred_sqls=pred_sqls,
            gold_dbs=gold_dbs,
            setup_sqls=setup_sqls,
            validate_sqls=validate_sqls,
            kmaps=kmaps,
            db_schemas=db_schemas,
            database_dir=db_dir,
            lowercase_schema_match=False,
            model_name=model_name,
            categories=categories,
        )

        # Save evaluation results to dataset
        save_evaluation(inference_api, model_name, prompt_format, custom_prompt, metrics)

        yield "Evaluation completed."

        if metrics:
            yield "Overall Results:"
            overall_metrics = metrics['exec']['all']
            yield f"All (n={overall_metrics['count']}) - Execution Accuracy: {overall_metrics['exec']:.3f}"
            yield f"All (n={overall_metrics['count']}) - Edit Distance: {metrics['edit_distance']['edit_distance']:.3f}"

            categories = ['easy', 'medium', 'hard', 'duckdb', 'ddl', 'all']

            for category in categories:
                if category in metrics['exec']:
                    category_metrics = metrics['exec'][category]
                    yield f"{category} (n={category_metrics['count']}) - Execution Accuracy: {category_metrics['exec']:.3f}"
                else:
                    yield f"{category}: No data available"
        else:
            yield "No evaluation metrics returned."
    except Exception as e:
        yield f"An unexpected error occurred: {str(e)}"
        yield f"Error traceback: {traceback.format_exc()}"

if __name__ == "__main__":
    model_name = input("Enter the model name: ")
    prompt_format = input("Enter the prompt format (default is duckdbinstgraniteshort): ") or "duckdbinstgraniteshort"
    for result in run_evaluation(model_name, prompt_format):
        print(result, flush=True)