Add text2sql tasks

.gitignore

@@ -159,3 +159,4 @@ benchmark_output/*
 .litellm_cache
 
 docs/_static/data.js
+cache

examples/evaluate_text2sql.py

@@ -0,0 +1,62 @@
+from unitxt import evaluate, load_dataset, settings
+from unitxt.inference import CrossProviderInferenceEngine
+from unitxt.text_utils import print_dict
+
+with settings.context(
+    disable_hf_datasets_cache=False,
+    allow_unverified_code=True,
+):
+    test_dataset = load_dataset(
+        "card=cards.text2sql.bird"
+        ",template=templates.text2sql.you_are_given_with_hint_with_sql_prefix,loader_limit=10",
+        # ",template=templates.text2sql.you_are_given_with_hint_with_sql_prefix",
+        split="validation",
+    )
+
+# Infer
+inference_model = CrossProviderInferenceEngine(
+    model="llama-3-70b-instruct",
+    max_tokens=256,
+)
+
+predictions = inference_model.infer(test_dataset)
+evaluated_dataset = evaluate(predictions=predictions, data=test_dataset)
+
+print_dict(
+    evaluated_dataset[0],
+    keys_to_print=[
+        "source",
+        "prediction",
+        "subset",
+    ],
+)
+print_dict(
+    evaluated_dataset[0]["score"]["global"],
+)
+
+# with llama-3-70b-instruct
+# num_of_instances (int):
+#     1534
+# execution_accuracy (float):
+#     0.482
+
+# like GPT4 (rank 40 in the benchmark https://bird-bench.github.io/)
+
+# from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# DEBUG_NUM_EXAMPLES = 2
+# model_name = "meta-llama/Llama-3.2-1B-Instruct"
+# model = AutoModelForCausalLM.from_pretrained(model_name)
+# tokenizer = AutoTokenizer.from_pretrained(model_name)
+# tokenizer.pad_token = tokenizer.eos_token
+# test_dataset = test_dataset.select(range(DEBUG_NUM_EXAMPLES))
+# predictions = tokenizer.batch_decode(
+#     model.generate(
+#         **tokenizer.batch_encode_plus(
+#             test_dataset["source"], return_tensors="pt", padding=True
+#         ),
+#         max_length=2048,
+#     ),
+#     skip_special_tokens=True,
+#     clean_up_tokenization_spaces=True,
+# )

prepare/cards/text2sql.py

@@ -0,0 +1,54 @@
+import sys
+
+from unitxt import add_to_catalog
+from unitxt.blocks import Copy, Rename, Set, TaskCard
+from unitxt.loaders import LoadHF
+from unitxt.operators import ExecuteExpression, Shuffle
+
+card = TaskCard(
+    loader=LoadHF(path="premai-io/birdbench", split="validation"),
+    preprocess_steps=[
+        Shuffle(page_size=sys.maxsize),
+        Rename(
+            field_to_field={
+                "question_id": "id",
+                "question": "utterance",
+                "SQL": "query",
+                "db_id": "db_id",
+                "evidence": "hint",
+            }
+        ),
+        Set(
+            fields={
+                "dbms": "sqlite",
+                "db_type": "local",
+                "use_oracle_knowledge": True,
+                "num_table_rows_to_add": 0,
+                "data": None,
+            }
+        ),
+        ExecuteExpression(
+            expression="'bird/'+db_id",
+            to_field="db_id",
+        ),
+        Copy(field="db_id", to_field="db/db_id"),
+        Copy(field="db_type", to_field="db/db_type"),
+        Copy(field="dbms", to_field="db/dbms"),
+        Copy(field="data", to_field="db/data"),
+    ],
+    task="tasks.text2sql",
+    templates="templates.text2sql.all",
+)
+
+# test_card(card, num_demos=0, demos_pool_size=0, )
+
+add_to_catalog(
+    card,
+    "cards.text2sql.bird",
+    overwrite=True,
+)
+
+# from unitxt import evaluate, load_dataset
+
+# ds = load_dataset("card=cards.text2sql.bird,template_card_index=0")
+# scores = evaluate(predictions=ds["validation"]["target"], data=ds["validation"])

prepare/metrics/text2sql_execution_accuracy.py

@@ -0,0 +1,65 @@
+from unitxt.catalog import add_to_catalog
+from unitxt.metrics import ExecutionAccuracy
+from unitxt.test_utils.metrics import test_metric
+
+metric = ExecutionAccuracy()
+
+predictions = [
+    "SELECT nme FROM employees WHERE department = 'Sales'",
+    "SELECT name FROM employees WHERE department = 'Sales'",
+]  # Incorrect column name 'nme'
+references = [["SELECT name FROM employees WHERE department = 'Sales';"]] * 2
+task_data = [
+    {
+        "db": {
+            "db_id": "mock_db",
+            "db_type": "in_memory",
+            "data": {
+                "employees": {
+                    "columns": ["id", "name", "department", "salary"],
+                    "rows": [
+                        (1, "Alice", "Sales", 50000),
+                        (2, "Bob", "Engineering", 60000),
+                        (3, "Charlie", "Sales", 55000),
+                    ],
+                }
+            },
+        }
+    }
+] * 2
+
+instance_targets = [
+    {
+        "execution_accuracy": 0.0,
+        "score": 0.0,
+        "score_name": "execution_accuracy",
+    },
+    {
+        "execution_accuracy": 1.0,
+        "score": 1.0,
+        "score_name": "execution_accuracy",
+    },
+]
+
+
+global_target = {
+    "execution_accuracy": 0.5,
+    "execution_accuracy_ci_high": 1.0,
+    "execution_accuracy_ci_low": 0.0,
+    "num_of_instances": 2,
+    "score": 0.5,
+    "score_ci_high": 1.0,
+    "score_ci_low": 0.0,
+    "score_name": "execution_accuracy",
+}
+
+outputs = test_metric(
+    metric=metric,
+    predictions=predictions,
+    references=references,
+    instance_targets=instance_targets,
+    global_target=global_target,
+    task_data=task_data,
+)
+
+add_to_catalog(metric, "metrics.text2sql.execution_accuracy", overwrite=True)

prepare/processors/text2sql.py

@@ -0,0 +1,13 @@
+from unitxt import add_to_catalog
+from unitxt.operator import SequentialOperator
+from unitxt.processors import GetSQL
+
+add_to_catalog(
+    SequentialOperator(
+        steps=[
+            GetSQL(field="prediction"),
+        ]
+    ),
+    "processors.text2sql.get_sql",
+    overwrite=True,
+)

prepare/serializers/text2sql_serializers.py

@@ -0,0 +1,6 @@
+from unitxt import add_to_catalog
+from unitxt.serializers import SQLDatabaseAsSchemaSerializer
+
+add_to_catalog(
+    SQLDatabaseAsSchemaSerializer(), "serializers.text2sql.schema", overwrite=True
+)

prepare/tasks/text2sql.py

@@ -0,0 +1,19 @@
+from unitxt.blocks import Task
+from unitxt.catalog import add_to_catalog
+from unitxt.types import SQLDatabase
+
+add_to_catalog(
+    Task(
+        input_fields={
+            "id": str,
+            "utterance": str,
+            "hint": str,
+            "db": SQLDatabase,
+        },
+        reference_fields={"query": str},
+        prediction_type=str,
+        metrics=["metrics.text2sql.execution_accuracy", "metrics.anls"],
+    ),
+    "tasks.text2sql",
+    overwrite=True,
+)

prepare/templates/text2sql/templates.py

@@ -0,0 +1,64 @@
+from unitxt import add_to_catalog
+from unitxt.blocks import TemplatesList
+from unitxt.templates import InputOutputTemplate
+
+template_details = [
+    (
+        "templates.text2sql.you_are_given_with_sql_prefix",
+        "You are given the following question:\n\n{utterance}\n\nAn SQL schema\n\n```sql\n\n{db}\n```\n\nAnswer the following question:\n\n{utterance}\n\n",
+        "You are a Text2SQL generation model, in your answer, only have SQL code.\nStart your query with 'SELECT' and end it with ';'\n\n",
+        "```sql\nSELECT ",
+    ),
+    (
+        "templates.text2sql.you_are_given_with_hint_with_sql_prefix",
+        "You are given the following question:\n\n{utterance}\n\nAn SQL schema\n\n```sql\n\n{db}\n```\n\nAnd hint:\n\n{hint}\n\nAnswer the following question:\n\n{utterance}\n\n",
+        "You are a Text2SQL generation model, in your answer, only have SQL code.\nMake sure you start your query with 'SELECT' and end it with ';'\n\n",
+        "```sql\nSELECT ",
+    ),
+    (
+        "templates.text2sql.you_are_given",
+        "You are given the following question:\n\n{utterance}\n\nAn SQL schema\n\n```sql\n\n{db}\n```\n\nAnswer the following question:\n\n{utterance}\n\n",
+        "You are a Text2SQL generation model, in your answer, only have SQL code.\nStart your query with 'SELECT' and end it with ';'\n\n",
+        "",
+    ),
+    (
+        "templates.text2sql.you_are_given_with_hint_with_sql_prefix",
+        "You are given the following question:\n\n{utterance}\n\nAn SQL schema\n\n```sql\n\n{db}\n```\n\nAnd hint:\n\n{hint}\n\nAnswer the following question:\n\n{utterance}\n\n",
+        "You are a Text2SQL generation model, in your answer, only have SQL code.\nMake sure you start your query with 'SELECT' and end it with ';'\n\n",
+        "",
+    ),
+    (
+        "templates.text2sql.you_are_given_with_hint_answer_sql_prefix_no_inst",
+        "Question:\nYou are given the following SQL schema\n\n```sql\n{db}\n```\n\n{utterance}\n\n",
+        "",
+        "Answer:\n```sql\n",
+    ),
+    (
+        "templates.text2sql.empty",
+        "{utterance}",
+        "",
+        "",
+    ),
+]
+
+template_names = []
+for name, input_format, instruction, target_prefix in template_details:
+    add_to_catalog(
+        InputOutputTemplate(
+            input_format=input_format,
+            instruction=instruction,
+            target_prefix=target_prefix,
+            output_format="{query}",
+            postprocessors=["processors.text2sql.get_sql"],
+        ),
+        name,
+        overwrite=True,
+    )
+    template_names.append(name)
+
+
+add_to_catalog(
+    TemplatesList(template_names),
+    "templates.text2sql.all",
+    overwrite=True,
+)

pyproject.toml

@@ -104,7 +104,8 @@ tests = [
     "bs4",
     "tenacity==8.3.0",
     "accelerate",
-    "spacy",
+    "spacy", 
+    "func_timeout==4.3.5",
     "Wikipedia-API"
 ]
 ui = [

src/unitxt/catalog/cards/text2sql/bird.json

@@ -0,0 +1,61 @@
+{
+    "__type__": "task_card",
+    "loader": {
+        "__type__": "load_hf",
+        "path": "premai-io/birdbench",
+        "split": "validation"
+    },
+    "preprocess_steps": [
+        {
+            "__type__": "shuffle",
+            "page_size": 9223372036854775807
+        },
+        {
+            "__type__": "rename",
+            "field_to_field": {
+                "question_id": "id",
+                "question": "utterance",
+                "SQL": "query",
+                "db_id": "db_id",
+                "evidence": "hint"
+            }
+        },
+        {
+            "__type__": "set",
+            "fields": {
+                "dbms": "sqlite",
+                "db_type": "local",
+                "use_oracle_knowledge": true,
+                "num_table_rows_to_add": 0,
+                "data": null
+            }
+        },
+        {
+            "__type__": "execute_expression",
+            "expression": "'bird/'+db_id",
+            "to_field": "db_id"
+        },
+        {
+            "__type__": "copy",
+            "field": "db_id",
+            "to_field": "db/db_id"
+        },
+        {
+            "__type__": "copy",
+            "field": "db_type",
+            "to_field": "db/db_type"
+        },
+        {
+            "__type__": "copy",
+            "field": "dbms",
+            "to_field": "db/dbms"
+        },
+        {
+            "__type__": "copy",
+            "field": "data",
+            "to_field": "db/data"
+        }
+    ],
+    "task": "tasks.text2sql",
+    "templates": "templates.text2sql.all"
+}

src/unitxt/catalog/metrics/text2sql/execution_accuracy.json

@@ -0,0 +1,3 @@
+{
+    "__type__": "execution_accuracy"
+}

src/unitxt/catalog/processors/text2sql/get_sql.json

@@ -0,0 +1,9 @@
+{
+    "__type__": "sequential_operator",
+    "steps": [
+        {
+            "__type__": "get_sql",
+            "field": "prediction"
+        }
+    ]
+}

src/unitxt/catalog/serializers/text2sql/schema.json

@@ -0,0 +1,3 @@
+{
+    "__type__": "sql_database_as_schema_serializer"
+}