def load_schema_obj(dataset_name: str, data_dir: str) -> Dict[Any, Any]:
if dataset_name.lower() == "spider":
return load_spider_tables(os.path.join(data_dir, "tables.json"))
elif dataset_name.lower() == "wikisql":
raise ValueError("WikiSQL inference is not supported yet")
else:
schema_csv = os.path.join(
data_dir,
dataset_name + "_schema.csv",
)
return read_schema(schema_csv)
def inference_wrapper(inference_fn, sharded=False):
"""Wrapper for running inference."""
dataset_name = FLAGS.dataset_name
if not FLAGS.predictions_path:
raise ValueError('Predictions path must be set.')
predictions = FLAGS.predictions_path + '*'
# Don't run inference if predictions have already been generated.
if not tf.gfile.Glob(FLAGS.predictions_path + '*'):
inference_fn(FLAGS.input, FLAGS.predictions_path,
FLAGS.checkpoint_filepath, dataset_name)
# If using Abstract SQL, need to restore under-specified FROM clauses
# output above.
if FLAGS.restore_preds_from_asql:
spider = dataset_name.lower() == 'spider'
if not FLAGS.restored_predictions_path:
raise ValueError('Restored predictions path must be set '
'if restoring predictions from AbSQL.')
if not tf.io.gfile.exists(FLAGS.restored_predictions_path):
restore_from_asql.restore_from_clauses(
predictions,
FLAGS.restored_predictions_path,
spider_examples_json=FLAGS.spider_examples_json
if spider else '',
spider_tables_json=FLAGS.spider_tables_json if spider else '',
michigan_schema=None if spider else read_schema(
os.path.join(FLAGS.data_filepath, FLAGS.dataset_name +
'_schema.csv')),
dataset_name=FLAGS.dataset_name,
use_oracle_foriegn_keys=FLAGS.use_oracle_foriegn_keys)
predictions = FLAGS.restored_predictions_path
if FLAGS.match_and_save:
# Load the database tables.
schema_obj = None
if dataset_name.lower() == 'spider':
schema_obj = load_spider_tables(
os.path.join(FLAGS.data_filepath, 'tables.json'))
elif dataset_name.lower() == 'wikisql':
raise ValueError('WikiSQL inference is not supported yet')
else:
schema_csv = os.path.join(FLAGS.data_filepath,
FLAGS.dataset_name + '_schema.csv')
schema_obj = read_schema(schema_csv)
# Now match with the original data and save
match_and_save(predictions, FLAGS.output, dataset_name.lower(),
FLAGS.splits, FLAGS.data_filepath, schema_obj, sharded)
def process_spider(output_file, debugging_file, tokenizer):
"""Loads, converts, and writes Spider examples to the standard format."""
if len(FLAGS.splits) > 1:
raise ValueError('Not expecting more than one split for Spider.')
split = FLAGS.splits[0]
table_definitions = load_spider_tables(
os.path.join(FLAGS.input_filepath, 'tables.json'))
print('Loaded %d table definitions.' % len(table_definitions))
spider_examples = \
load_spider_examples(os.path.join(FLAGS.input_filepath,
split + '.json'))
num_examples_created = 0
num_examples_failed = 0
# TODO(petershaw): Reduce duplication with other code path for schema
# pre-processing.
tables_json = _load_json_from_file(
os.path.join(FLAGS.input_filepath, 'tables.json'))
spider_table_schemas_map = abstract_sql_converters.spider_table_schemas_map(
tables_json)
for spider_example in spider_examples:
# Make sure the DB specified exists.
example_db = spider_example['db_id']
try:
example = convert_spider(
spider_example,
table_definitions[example_db],
tokenizer,
FLAGS.generate_sql,
FLAGS.anonymize_values,
abstract_sql=FLAGS.abstract_sql,
table_schemas=spider_table_schemas_map[example_db],
allow_value_generation=FLAGS.allow_value_generation)
except abstract_sql.UnsupportedSqlError as e:
print(e)
example = None
if example:
output_file.write(json.dumps(example.to_json()) + '\n')
num_examples_created += 1
debugging_file.write(example.model_input.original_utterance + '\n')
if FLAGS.generate_sql:
debugging_file.write(example.gold_query_string() + '\n\n')
else:
num_examples_failed += 1
return num_examples_created, num_examples_failed