Merge branch 'main' into add-text2sql

IBM · Jan 22, 2025 · 4fc611f · 4fc611f
2 parents 9ae0061 + 52031bb
commit 4fc611f
Show file tree

Hide file tree

Showing 49 changed files with 168 additions and 80 deletions.
diff --git a/prepare/cards/open_australian_legal_qa.py b/prepare/cards/open_australian_legal_qa.py
@@ -13,6 +13,7 @@
 card = TaskCard(
     loader=LoadHF(
         path="umarbutler/open-australian-legal-qa",
+        name="default",
     ),
     preprocess_steps=[
         SplitRandomMix(
@@ -51,6 +52,7 @@
 card = TaskCard(
     loader=LoadHF(
         path="umarbutler/open-australian-legal-qa",
+        name="default",
     ),
     preprocess_steps=[
         SplitRandomMix(

diff --git a/prepare/recipes/tables_benchmark.py b/prepare/recipes/tables_benchmark.py
@@ -22,7 +22,6 @@
 serializers = ",".join(list(SERIALIZERS))
 max_augmentors = 10
 max_pred_tokens = 100
-num_demos = 5
 recipes_only = False
 
 # Process parameters
@@ -35,6 +34,7 @@
 for card in cards_parsed:
     for augment in all_augment:
         for serializer in serializers_parsed:
+            num_demos = 1 if card == "wikitq" else 5
             kwargs = {
                 "card": "cards." + card,
                 "serializer": f"serializers.table.{serializer}"

diff --git a/src/unitxt/assistant/app.py b/src/unitxt/assistant/app.py
@@ -1,5 +1,7 @@
 import datetime
+import importlib
 import json
+import logging
 import os
 import uuid
 
@@ -8,8 +10,11 @@
 import pandas as pd
 import streamlit as st
 import torch
+from litellm import AuthenticationError
 from transformers import AutoTokenizer
 
+logger = logging.getLogger("unitxt-assistance")
+
 
 @st.cache_resource
 def load_data():
@@ -20,9 +25,36 @@ def load_data():
     return metadata_df, embeddings
 
 
+def get_embedding_with_retry(model, input, max_retries=3):
+    """This function calls the litellm.embedding method and handles token expiration.
+
+    It will retry the call up to `max_retries` times if an AuthenticationError is raised.
+    """
+    retries = 0
+    actual_exception = None
+    while retries < max_retries:
+        try:
+            return litellm.embedding(model=model, input=input)
+
+        except AuthenticationError as e:
+            actual_exception = e
+            retries += 1
+            logger.info(
+                f"Authentication error: {e}. Retrying... ({retries}/{max_retries})"
+            )
+            importlib.reload(
+                litellm
+            )  # Reload the litellm module to clear any cached state
+
+    # If all retries fail, raise an error
+    raise Exception(
+        f"Failed to get embedding after {max_retries} attempts. Exception: {actual_exception}"
+    )
+
+
 def search(query, metadata_df, embeddings, max_tokens=5000, min_text_length=50):
     # Generate embedding for the query using litellm
-    response = litellm.embedding(
+    response = get_embedding_with_retry(
         model="watsonx/intfloat/multilingual-e5-large",
         input=[query],
     )

diff --git a/src/unitxt/catalog/benchmarks/tables_benchmark.json b/src/unitxt/catalog/benchmarks/tables_benchmark.json
@@ -683,71 +683,71 @@
                 "concat": {
                     "__type__": "benchmark",
                     "subsets": {
-                        "insert_empty_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.concat.insert_empty_rows_augmentation_5_demos",
-                        "no_augmentation_5_demos": "recipes.tables_benchmark.wikitq.concat.no_augmentation_5_demos",
-                        "shuffle_cols_augmentation_5_demos": "recipes.tables_benchmark.wikitq.concat.shuffle_cols_augmentation_5_demos",
-                        "shuffle_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.concat.shuffle_rows_augmentation_5_demos",
-                        "transpose_augmentation_5_demos": "recipes.tables_benchmark.wikitq.concat.transpose_augmentation_5_demos"
+                        "insert_empty_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.concat.insert_empty_rows_augmentation_1_demos",
+                        "no_augmentation_1_demos": "recipes.tables_benchmark.wikitq.concat.no_augmentation_1_demos",
+                        "shuffle_cols_augmentation_1_demos": "recipes.tables_benchmark.wikitq.concat.shuffle_cols_augmentation_1_demos",
+                        "shuffle_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.concat.shuffle_rows_augmentation_1_demos",
+                        "transpose_augmentation_1_demos": "recipes.tables_benchmark.wikitq.concat.transpose_augmentation_1_demos"
                     }
                 },
                 "csv": {
                     "__type__": "benchmark",
                     "subsets": {
-                        "insert_empty_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.csv.insert_empty_rows_augmentation_5_demos",
-                        "no_augmentation_5_demos": "recipes.tables_benchmark.wikitq.csv.no_augmentation_5_demos",
-                        "shuffle_cols_augmentation_5_demos": "recipes.tables_benchmark.wikitq.csv.shuffle_cols_augmentation_5_demos",
-                        "shuffle_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.csv.shuffle_rows_augmentation_5_demos",
-                        "transpose_augmentation_5_demos": "recipes.tables_benchmark.wikitq.csv.transpose_augmentation_5_demos"
+                        "insert_empty_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.csv.insert_empty_rows_augmentation_1_demos",
+                        "no_augmentation_1_demos": "recipes.tables_benchmark.wikitq.csv.no_augmentation_1_demos",
+                        "shuffle_cols_augmentation_1_demos": "recipes.tables_benchmark.wikitq.csv.shuffle_cols_augmentation_1_demos",
+                        "shuffle_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.csv.shuffle_rows_augmentation_1_demos",
+                        "transpose_augmentation_1_demos": "recipes.tables_benchmark.wikitq.csv.transpose_augmentation_1_demos"
                     }
                 },
                 "df": {
                     "__type__": "benchmark",
                     "subsets": {
-                        "insert_empty_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.df.insert_empty_rows_augmentation_5_demos",
-                        "no_augmentation_5_demos": "recipes.tables_benchmark.wikitq.df.no_augmentation_5_demos",
-                        "shuffle_cols_augmentation_5_demos": "recipes.tables_benchmark.wikitq.df.shuffle_cols_augmentation_5_demos",
-                        "shuffle_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.df.shuffle_rows_augmentation_5_demos",
-                        "transpose_augmentation_5_demos": "recipes.tables_benchmark.wikitq.df.transpose_augmentation_5_demos"
+                        "insert_empty_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.df.insert_empty_rows_augmentation_1_demos",
+                        "no_augmentation_1_demos": "recipes.tables_benchmark.wikitq.df.no_augmentation_1_demos",
+                        "shuffle_cols_augmentation_1_demos": "recipes.tables_benchmark.wikitq.df.shuffle_cols_augmentation_1_demos",
+                        "shuffle_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.df.shuffle_rows_augmentation_1_demos",
+                        "transpose_augmentation_1_demos": "recipes.tables_benchmark.wikitq.df.transpose_augmentation_1_demos"
                     }
                 },
                 "html": {
                     "__type__": "benchmark",
                     "subsets": {
-                        "insert_empty_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.html.insert_empty_rows_augmentation_5_demos",
-                        "no_augmentation_5_demos": "recipes.tables_benchmark.wikitq.html.no_augmentation_5_demos",
-                        "shuffle_cols_augmentation_5_demos": "recipes.tables_benchmark.wikitq.html.shuffle_cols_augmentation_5_demos",
-                        "shuffle_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.html.shuffle_rows_augmentation_5_demos",
-                        "transpose_augmentation_5_demos": "recipes.tables_benchmark.wikitq.html.transpose_augmentation_5_demos"
+                        "insert_empty_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.html.insert_empty_rows_augmentation_1_demos",
+                        "no_augmentation_1_demos": "recipes.tables_benchmark.wikitq.html.no_augmentation_1_demos",
+                        "shuffle_cols_augmentation_1_demos": "recipes.tables_benchmark.wikitq.html.shuffle_cols_augmentation_1_demos",
+                        "shuffle_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.html.shuffle_rows_augmentation_1_demos",
+                        "transpose_augmentation_1_demos": "recipes.tables_benchmark.wikitq.html.transpose_augmentation_1_demos"
                     }
                 },
                 "indexed_row_major": {
                     "__type__": "benchmark",
                     "subsets": {
-                        "insert_empty_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.insert_empty_rows_augmentation_5_demos",
-                        "no_augmentation_5_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.no_augmentation_5_demos",
-                        "shuffle_cols_augmentation_5_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.shuffle_cols_augmentation_5_demos",
-                        "shuffle_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.shuffle_rows_augmentation_5_demos",
-                        "transpose_augmentation_5_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.transpose_augmentation_5_demos"
+                        "insert_empty_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.insert_empty_rows_augmentation_1_demos",
+                        "no_augmentation_1_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.no_augmentation_1_demos",
+                        "shuffle_cols_augmentation_1_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.shuffle_cols_augmentation_1_demos",
+                        "shuffle_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.shuffle_rows_augmentation_1_demos",
+                        "transpose_augmentation_1_demos": "recipes.tables_benchmark.wikitq.indexed_row_major.transpose_augmentation_1_demos"
                     }
                 },
                 "json": {
                     "__type__": "benchmark",
                     "subsets": {
-                        "insert_empty_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.json.insert_empty_rows_augmentation_5_demos",
-                        "no_augmentation_5_demos": "recipes.tables_benchmark.wikitq.json.no_augmentation_5_demos",
-                        "shuffle_cols_augmentation_5_demos": "recipes.tables_benchmark.wikitq.json.shuffle_cols_augmentation_5_demos",
-                        "shuffle_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.json.shuffle_rows_augmentation_5_demos",
-                        "transpose_augmentation_5_demos": "recipes.tables_benchmark.wikitq.json.transpose_augmentation_5_demos"
+                        "insert_empty_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.json.insert_empty_rows_augmentation_1_demos",
+                        "no_augmentation_1_demos": "recipes.tables_benchmark.wikitq.json.no_augmentation_1_demos",
+                        "shuffle_cols_augmentation_1_demos": "recipes.tables_benchmark.wikitq.json.shuffle_cols_augmentation_1_demos",
+                        "shuffle_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.json.shuffle_rows_augmentation_1_demos",
+                        "transpose_augmentation_1_demos": "recipes.tables_benchmark.wikitq.json.transpose_augmentation_1_demos"
                     }
                 },
                 "markdown": {
                     "__type__": "benchmark",
                     "subsets": {
-                        "insert_empty_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.markdown.insert_empty_rows_augmentation_5_demos",
-                        "no_augmentation_5_demos": "recipes.tables_benchmark.wikitq.markdown.no_augmentation_5_demos",
-                        "shuffle_cols_augmentation_5_demos": "recipes.tables_benchmark.wikitq.markdown.shuffle_cols_augmentation_5_demos",
-                        "shuffle_rows_augmentation_5_demos": "recipes.tables_benchmark.wikitq.markdown.shuffle_rows_augmentation_5_demos",
-                        "transpose_augmentation_5_demos": "recipes.tables_benchmark.wikitq.markdown.transpose_augmentation_5_demos"
+                        "insert_empty_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.markdown.insert_empty_rows_augmentation_1_demos",
+                        "no_augmentation_1_demos": "recipes.tables_benchmark.wikitq.markdown.no_augmentation_1_demos",
+                        "shuffle_cols_augmentation_1_demos": "recipes.tables_benchmark.wikitq.markdown.shuffle_cols_augmentation_1_demos",
+                        "shuffle_rows_augmentation_1_demos": "recipes.tables_benchmark.wikitq.markdown.shuffle_rows_augmentation_1_demos",
+                        "transpose_augmentation_1_demos": "recipes.tables_benchmark.wikitq.markdown.transpose_augmentation_1_demos"
                     }
                 }
             }

diff --git a/src/unitxt/catalog/cards/open_australian_legal_qa.json b/src/unitxt/catalog/cards/open_australian_legal_qa.json
@@ -2,7 +2,8 @@
     "__type__": "task_card",
     "loader": {
         "__type__": "load_hf",
-        "path": "umarbutler/open-australian-legal-qa"
+        "path": "umarbutler/open-australian-legal-qa",
+        "name": "default"
     },
     "preprocess_steps": [
         {

diff --git a/src/unitxt/catalog/cards/rag/response_generation/train/open_australian_legal_qa.json b/src/unitxt/catalog/cards/rag/response_generation/train/open_australian_legal_qa.json
@@ -2,7 +2,8 @@
     "__type__": "task_card",
     "loader": {
         "__type__": "load_hf",
-        "path": "umarbutler/open-australian-legal-qa"
+        "path": "umarbutler/open-australian-legal-qa",
+        "name": "default"
     },
     "preprocess_steps": [
         {

diff --git a/src/unitxt/catalog/metrics/llm_as_judge/direct/rits/granite3_0_8b.json b/src/unitxt/catalog/metrics/llm_as_judge/direct/rits/granite3_0_8b.json
@@ -0,0 +1,11 @@
+{
+    "__type__": "llm_judge_direct",
+    "inference_engine": {
+        "__type__": "rits_inference_engine",
+        "model_name": "ibm-granite/granite-3.0-8b-instruct",
+        "max_tokens": 1024,
+        "seed": 42
+    },
+    "evaluator_name": "GRANITE3_8B",
+    "generate_summaries": false
+}
diff --git a/src/unitxt/catalog/metrics/llm_as_judge/direct/rits/granite3_1_8b.json b/src/unitxt/catalog/metrics/llm_as_judge/direct/rits/granite3_1_8b.json
@@ -0,0 +1,11 @@
+{
+    "__type__": "llm_judge_direct",
+    "inference_engine": {
+        "__type__": "rits_inference_engine",
+        "model_name": "ibm-granite/granite-3.1-8b-instruct",
+        "max_tokens": 1024,
+        "seed": 42
+    },
+    "evaluator_name": "GRANITE3_1_8B",
+    "generate_summaries": false
+}
diff --git a/..._as_judge/direct/watsonx/granite3_8b.json → ...s_judge/direct/watsonx/granite3_0_8b.json b/..._as_judge/direct/watsonx/granite3_8b.json → ...s_judge/direct/watsonx/granite3_0_8b.json
diff --git a/src/unitxt/catalog/metrics/llm_as_judge/pairwise/rits/granite3_0_8b.json b/src/unitxt/catalog/metrics/llm_as_judge/pairwise/rits/granite3_0_8b.json
@@ -0,0 +1,11 @@
+{
+    "__type__": "llm_judge_pairwise",
+    "inference_engine": {
+        "__type__": "rits_inference_engine",
+        "model_name": "ibm-granite/granite-3.0-8b-instruct",
+        "max_tokens": 1024,
+        "seed": 42
+    },
+    "evaluator_name": "GRANITE3_8B",
+    "generate_summaries": false
+}
diff --git a/src/unitxt/catalog/metrics/llm_as_judge/pairwise/rits/granite3_1_8b.json b/src/unitxt/catalog/metrics/llm_as_judge/pairwise/rits/granite3_1_8b.json
@@ -0,0 +1,11 @@
+{
+    "__type__": "llm_judge_pairwise",
+    "inference_engine": {
+        "__type__": "rits_inference_engine",
+        "model_name": "ibm-granite/granite-3.1-8b-instruct",
+        "max_tokens": 1024,
+        "seed": 42
+    },
+    "evaluator_name": "GRANITE3_1_8B",
+    "generate_summaries": false
+}
diff --git a/...s_judge/pairwise/watsonx/granite3_8b.json → ...judge/pairwise/watsonx/granite3_0_8b.json b/...s_judge/pairwise/watsonx/granite3_8b.json → ...judge/pairwise/watsonx/granite3_0_8b.json
diff --git a/...sert_empty_rows_augmentation_5_demos.json → ...sert_empty_rows_augmentation_1_demos.json b/...sert_empty_rows_augmentation_5_demos.json → ...sert_empty_rows_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.concat",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...ikitq/concat/no_augmentation_5_demos.json → ...ikitq/concat/no_augmentation_1_demos.json b/...ikitq/concat/no_augmentation_5_demos.json → ...ikitq/concat/no_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.concat",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": null

diff --git a/...at/shuffle_cols_augmentation_5_demos.json → ...at/shuffle_cols_augmentation_1_demos.json b/...at/shuffle_cols_augmentation_5_demos.json → ...at/shuffle_cols_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.concat",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...at/shuffle_rows_augmentation_5_demos.json → ...at/shuffle_rows_augmentation_1_demos.json b/...at/shuffle_rows_augmentation_5_demos.json → ...at/shuffle_rows_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.concat",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...oncat/transpose_augmentation_5_demos.json → ...oncat/transpose_augmentation_1_demos.json b/...oncat/transpose_augmentation_5_demos.json → ...oncat/transpose_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.concat",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...sert_empty_rows_augmentation_5_demos.json → ...sert_empty_rows_augmentation_1_demos.json b/...sert_empty_rows_augmentation_5_demos.json → ...sert_empty_rows_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": null,
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...k/wikitq/csv/no_augmentation_5_demos.json → ...k/wikitq/csv/no_augmentation_1_demos.json b/...k/wikitq/csv/no_augmentation_5_demos.json → ...k/wikitq/csv/no_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": null,
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": null

diff --git a/...sv/shuffle_cols_augmentation_5_demos.json → ...sv/shuffle_cols_augmentation_1_demos.json b/...sv/shuffle_cols_augmentation_5_demos.json → ...sv/shuffle_cols_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": null,
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...sv/shuffle_rows_augmentation_5_demos.json → ...sv/shuffle_rows_augmentation_1_demos.json b/...sv/shuffle_rows_augmentation_5_demos.json → ...sv/shuffle_rows_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": null,
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...q/csv/transpose_augmentation_5_demos.json → ...q/csv/transpose_augmentation_1_demos.json b/...q/csv/transpose_augmentation_5_demos.json → ...q/csv/transpose_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": null,
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...sert_empty_rows_augmentation_5_demos.json → ...sert_empty_rows_augmentation_1_demos.json b/...sert_empty_rows_augmentation_5_demos.json → ...sert_empty_rows_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.df",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...rk/wikitq/df/no_augmentation_5_demos.json → ...rk/wikitq/df/no_augmentation_1_demos.json b/...rk/wikitq/df/no_augmentation_5_demos.json → ...rk/wikitq/df/no_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.df",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": null

diff --git a/...df/shuffle_cols_augmentation_5_demos.json → ...df/shuffle_cols_augmentation_1_demos.json b/...df/shuffle_cols_augmentation_5_demos.json → ...df/shuffle_cols_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.df",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...df/shuffle_rows_augmentation_5_demos.json → ...df/shuffle_rows_augmentation_1_demos.json b/...df/shuffle_rows_augmentation_5_demos.json → ...df/shuffle_rows_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.df",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [

diff --git a/...tq/df/transpose_augmentation_5_demos.json → ...tq/df/transpose_augmentation_1_demos.json b/...tq/df/transpose_augmentation_5_demos.json → ...tq/df/transpose_augmentation_1_demos.json
@@ -2,7 +2,7 @@
     "__type__": "dataset_recipe",
     "card": "cards.wikitq",
     "serializer": "serializers.table.df",
-    "num_demos": 5,
+    "num_demos": 1,
     "demos_pool_size": -1,
     "loader_limit": 10000,
     "augmentor": [