448 lines
16 KiB
Python
448 lines
16 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import os
|
|
import random
|
|
import re
|
|
from typing import Any, Dict, List
|
|
|
|
import torch
|
|
|
|
from modelscope.utils.constant import ModeKeys
|
|
from .base import OfaBasePreprocessor
|
|
from .utils.bridge_content_encoder import get_database_matches
|
|
from .utils.get_tables import dump_db_json_schema
|
|
|
|
|
|
class OfaTextToSqlPreprocessor(OfaBasePreprocessor):
|
|
r"""
|
|
OFA preprocessor for text to sql tasks
|
|
"""
|
|
|
|
def __init__(self,
|
|
cfg,
|
|
model_dir,
|
|
mode=ModeKeys.INFERENCE,
|
|
*args,
|
|
**kwargs):
|
|
"""preprocess the data
|
|
|
|
Args:
|
|
cfg(modelscope.utils.config.ConfigDict) : model config
|
|
model_dir (str): model path,
|
|
mode: preprocessor mode (model mode)
|
|
"""
|
|
super(OfaTextToSqlPreprocessor, self).__init__(cfg, model_dir, mode,
|
|
*args, **kwargs)
|
|
|
|
self.instruction_text = self.cfg.model.get('prompt',
|
|
' . generating sql code.')
|
|
self.max_struct_length = self.cfg.get('max_struct_length', 256)
|
|
self.separator = '\t'
|
|
self.db_schema_cache = {}
|
|
self.database_path = os.path.join(
|
|
os.path.abspath(model_dir), 'database')
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
if self.mode == ModeKeys.TRAIN:
|
|
return self._build_train_sample(data)
|
|
else:
|
|
return self._build_infer_sample(data)
|
|
|
|
def _build_train_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
r"""
|
|
build sample for training tasks.
|
|
|
|
step 1. Get the input question and database id from text input
|
|
step 2. Get the database structure input
|
|
step 3. Add a pseudo ids for every input.
|
|
step 4. Calculate the target and previous output items.
|
|
"""
|
|
assert 'text' in self.column_map and 'text' in data, \
|
|
'there must be `text` column in task key map and source data'
|
|
text = data[self.column_map['text']] # equal data['text']
|
|
texts = text.split(self.separator)
|
|
assert len(
|
|
texts
|
|
) == 3, 'invalid input, should contain query, question and database id'
|
|
query, question, db_id = texts
|
|
|
|
# construct struct input
|
|
if db_id not in self.db_schema_cache:
|
|
self.db_schema_cache[db_id] = dump_db_json_schema(
|
|
self.database_path + '/' + db_id + '/' + db_id + '.sqlite',
|
|
db_id)
|
|
|
|
question = ' '.join(question.strip().split()[:self.max_src_length])
|
|
|
|
seq_inputs = seq2seq_input(query, question, db_id, self.database_path,
|
|
self.db_schema_cache[db_id], self.cfg.model,
|
|
True)
|
|
struct_in = seq_inputs['struct_in']
|
|
text = seq_inputs['text_in']
|
|
seq_out = seq_inputs['seq_out']
|
|
db_struct = seq_inputs['db_struct']
|
|
|
|
text = '{} ; structured knowledge: {}'.format(
|
|
text, struct_in) + self.instruction_text
|
|
src_item = self.tokenize_text(text + self.instruction_text)
|
|
src_item = src_item[:(self.max_src_length + self.max_struct_length
|
|
+ 20)]
|
|
|
|
tgt_item = self.tokenize_text(
|
|
' {}'.format(seq_out), add_bos=False,
|
|
add_eos=False)[:self.max_tgt_length]
|
|
target_item = torch.cat([tgt_item, self.eos_item])
|
|
prev_output_item = torch.cat([self.bos_item, tgt_item])
|
|
|
|
sample = {
|
|
'id': 0.0,
|
|
'source': src_item,
|
|
'target': target_item,
|
|
'prev_output_tokens': prev_output_item,
|
|
'db_struct': db_struct
|
|
}
|
|
|
|
return sample
|
|
|
|
def _build_infer_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
|
|
r"""
|
|
build sample for inference tasks.
|
|
|
|
step 1. Get the input question and database id from text input
|
|
step 2. Get the database structure input
|
|
step 3. Add a pseudo ids for every input.
|
|
"""
|
|
assert 'text' in self.column_map and 'text' in data, \
|
|
'there must be `text` column in task key map and source data'
|
|
text = data[self.column_map['text']] # equal data['text']
|
|
db_id = data.get(self.column_map['database'], 'culture_company')
|
|
db_id = db_id.strip()
|
|
|
|
# construct struct input
|
|
if db_id not in self.db_schema_cache:
|
|
self.db_schema_cache[db_id] = dump_db_json_schema(
|
|
self.database_path + '/' + db_id + '/' + db_id + '.sqlite',
|
|
db_id)
|
|
|
|
text = ' '.join(text.strip().split()[:self.max_src_length])
|
|
|
|
seq_inputs = seq2seq_input(None, text, db_id, self.database_path,
|
|
self.db_schema_cache[db_id], self.cfg.model)
|
|
struct_in = seq_inputs['struct_in']
|
|
db_struct = seq_inputs['db_struct']
|
|
text = '{} ; structured knowledge: {}'.format(
|
|
text, struct_in) + self.instruction_text
|
|
src_item = self.tokenize_text(text + self.instruction_text)
|
|
src_item = src_item[:(self.max_src_length + self.max_struct_length
|
|
+ 20)]
|
|
|
|
sample = {'id': 0.0, 'source': src_item, 'db_struct': db_struct}
|
|
|
|
if 'solution' in self.column_map and self.column_map[
|
|
'solution'] in data:
|
|
sample['label'] = ' {}'.format(data[self.column_map['solution']])
|
|
return sample
|
|
|
|
|
|
def seq2seq_input(query,
|
|
question,
|
|
db_id,
|
|
db_path,
|
|
schema,
|
|
args,
|
|
is_train=False):
|
|
ex = form_input_for_construction(query, question, db_id, db_path, schema)
|
|
serialized_schema = spider_add_serialized_schema(
|
|
ex, args)['serialized_schema'].strip()
|
|
if not is_train:
|
|
return {
|
|
'struct_in': serialized_schema,
|
|
'text_in': question,
|
|
'db_struct': ex
|
|
}
|
|
question, seq_out = spider_pre_process_one_function(ex, args)
|
|
return {
|
|
'struct_in': serialized_schema,
|
|
'text_in': question,
|
|
'seq_out': seq_out,
|
|
'db_struct': ex
|
|
}
|
|
|
|
|
|
def spider_pre_process_one_function(item: dict, args):
|
|
prefix = ''
|
|
|
|
seq_out = spider_get_target(
|
|
query=item['query'],
|
|
db_id=item['db_id'],
|
|
normalize_query=True,
|
|
target_with_db_id=args.target_with_db_id,
|
|
)
|
|
|
|
return prefix + item['question'].strip(), seq_out
|
|
|
|
|
|
def spider_get_target(
|
|
query: str,
|
|
db_id: str,
|
|
normalize_query: bool,
|
|
target_with_db_id: bool,
|
|
) -> str:
|
|
_normalize = normalize if normalize_query else (lambda x: x)
|
|
return f'{db_id} | {_normalize(query)}' if target_with_db_id else _normalize(
|
|
query)
|
|
|
|
|
|
def normalize(query: str) -> str:
|
|
|
|
def comma_fix(s):
|
|
# Remove spaces in front of commas
|
|
return s.replace(' , ', ', ')
|
|
|
|
def white_space_fix(s):
|
|
# Remove double and triple spaces
|
|
return ' '.join(s.split())
|
|
|
|
def lower(s):
|
|
# Convert everything except text between (single or double) quotation marks to lower case
|
|
return re.sub(r"\b(?<!['\"])(\w+)(?!['\"])\b",
|
|
lambda match: match.group(1).lower(), s)
|
|
|
|
return comma_fix(white_space_fix(lower(query)))
|
|
|
|
|
|
def spider_add_serialized_schema(ex: dict, args) -> dict:
|
|
if getattr(args, 'schema_serialization_with_nl'):
|
|
serialized_schema = serialize_schema_natural_language(
|
|
question=ex['question'],
|
|
db_path=ex['db_path'],
|
|
db_id=ex['db_id'],
|
|
db_column_names=ex['db_column_names'],
|
|
db_table_names=ex['db_table_names'],
|
|
db_primary_keys=ex['db_primary_keys'],
|
|
db_foreign_keys=ex['db_foreign_keys'],
|
|
schema_serialization_with_db_content=args.
|
|
schema_serialization_with_db_content,
|
|
normalize_query=True,
|
|
)
|
|
else:
|
|
serialized_schema = serialize_schema(
|
|
question=ex['question'],
|
|
db_path=ex['db_path'],
|
|
db_id=ex['db_id'],
|
|
db_column_names=ex['db_column_names'],
|
|
db_table_names=ex['db_table_names'],
|
|
schema_serialization_type='peteshaw',
|
|
schema_serialization_randomized=False,
|
|
schema_serialization_with_db_id=True,
|
|
schema_serialization_with_db_content=args.
|
|
schema_serialization_with_db_content,
|
|
normalize_query=True,
|
|
)
|
|
return {'serialized_schema': serialized_schema}
|
|
|
|
|
|
def serialize_schema_natural_language(
|
|
question: str,
|
|
db_path: str,
|
|
db_id: str,
|
|
db_column_names: Dict[str, str],
|
|
db_table_names: List[str],
|
|
db_primary_keys,
|
|
db_foreign_keys,
|
|
schema_serialization_with_db_content: bool = False,
|
|
normalize_query: bool = True,
|
|
) -> str:
|
|
overall_description = f'{db_id} contains tables such as ' \
|
|
f'{", ".join([name.lower() if normalize_query else name for name in db_table_names])}.'
|
|
|
|
def table_description_primary_key_template(primary_key):
|
|
return f'{primary_key} is the primary key.'
|
|
|
|
def table_description(name, column_names):
|
|
return f'Table {name} has columns such as {", ".join(column_names)}.'
|
|
|
|
def value_description(cv_pairs):
|
|
return f'{"".join(["The {} contains values such as {}.".format(column, value) for column, value in cv_pairs])}'
|
|
|
|
def foreign_key_description(table_1, column_1, table_2, column_2):
|
|
return f'The {column_1} of {table_1} is the foreign key of {column_2} of {table_2}.'
|
|
|
|
db_primary_keys = db_primary_keys['column_id']
|
|
db_foreign_keys = list(
|
|
zip(db_foreign_keys['column_id'], db_foreign_keys['other_column_id']))
|
|
|
|
descriptions = [overall_description]
|
|
db_table_name_strs = []
|
|
db_column_name_strs = []
|
|
value_sep = ', '
|
|
for table_id, table_name in enumerate(db_table_names):
|
|
table_name_str = table_name.lower() if normalize_query else table_name
|
|
db_table_name_strs.append(table_name_str)
|
|
columns = []
|
|
column_value_pairs = []
|
|
primary_keys = []
|
|
for column_id, (x, y) in enumerate(
|
|
zip(db_column_names['table_id'],
|
|
db_column_names['column_name'])):
|
|
if column_id == 0:
|
|
continue
|
|
column_str = y.lower() if normalize_query else y
|
|
db_column_name_strs.append(column_str)
|
|
if x == table_id:
|
|
columns.append(column_str)
|
|
if column_id in db_primary_keys:
|
|
primary_keys.append(column_str)
|
|
if schema_serialization_with_db_content:
|
|
matches = get_database_matches(
|
|
question=question,
|
|
table_name=table_name,
|
|
column_name=y,
|
|
db_path=(db_path + '/' + db_id + '/' + db_id
|
|
+ '.sqlite'),
|
|
)
|
|
if matches:
|
|
column_value_pairs.append(
|
|
(column_str, value_sep.join(matches)))
|
|
|
|
table_description_columns_str = table_description(
|
|
table_name_str, columns)
|
|
descriptions.append(table_description_columns_str)
|
|
table_description_primary_key_str = table_description_primary_key_template(
|
|
', '.join(primary_keys))
|
|
descriptions.append(table_description_primary_key_str)
|
|
if len(column_value_pairs) > 0:
|
|
value_description_str = value_description(column_value_pairs)
|
|
descriptions.append(value_description_str)
|
|
|
|
for x, y in db_foreign_keys:
|
|
# get the table and column of x
|
|
x_table_name = db_table_name_strs[db_column_names['table_id'][x]]
|
|
x_column_name = db_column_name_strs[x]
|
|
# get the table and column of y
|
|
y_table_name = db_table_name_strs[db_column_names['table_id'][y]]
|
|
y_column_name = db_column_name_strs[y]
|
|
foreign_key_description_str = foreign_key_description(
|
|
x_table_name, x_column_name, y_table_name, y_column_name)
|
|
descriptions.append(foreign_key_description_str)
|
|
return ' '.join(descriptions)
|
|
|
|
|
|
def serialize_schema(
|
|
question: str,
|
|
db_path: str,
|
|
db_id: str,
|
|
db_column_names: Dict[str, str],
|
|
db_table_names: List[str],
|
|
schema_serialization_type: str = 'peteshaw',
|
|
schema_serialization_randomized: bool = False,
|
|
schema_serialization_with_db_id: bool = True,
|
|
schema_serialization_with_db_content: bool = False,
|
|
normalize_query: bool = True,
|
|
) -> str:
|
|
if schema_serialization_type == 'verbose':
|
|
db_id_str = 'Database: {db_id}. '
|
|
table_sep = '. '
|
|
table_str = 'Table: {table}. Columns: {columns}'
|
|
column_sep = ', '
|
|
column_str_with_values = '{column} ({values})'
|
|
column_str_without_values = '{column}'
|
|
value_sep = ', '
|
|
elif schema_serialization_type == 'peteshaw':
|
|
# see https://github.com/google-research/language/blob/master/language/nqg/tasks/spider/append_schema.py#L42
|
|
db_id_str = ' | {db_id}'
|
|
table_sep = ''
|
|
table_str = ' | {table} : {columns}'
|
|
column_sep = ' , '
|
|
column_str_with_values = '{column} ( {values} )'
|
|
column_str_without_values = '{column}'
|
|
value_sep = ' , '
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def get_column_str(table_name: str, column_name: str) -> str:
|
|
column_name_str = column_name.lower(
|
|
) if normalize_query else column_name
|
|
if schema_serialization_with_db_content:
|
|
# print("testing")
|
|
matches = get_database_matches(
|
|
question=question,
|
|
table_name=table_name,
|
|
column_name=column_name,
|
|
db_path=(db_path + '/' + db_id + '/' + db_id + '.sqlite'),
|
|
)
|
|
if matches:
|
|
return column_str_with_values.format(
|
|
column=column_name_str, values=value_sep.join(matches))
|
|
else:
|
|
return column_str_without_values.format(column=column_name_str)
|
|
else:
|
|
return column_str_without_values.format(column=column_name_str)
|
|
|
|
tables = [
|
|
table_str.format(
|
|
table=table_name.lower() if normalize_query else table_name,
|
|
columns=column_sep.join(
|
|
map(
|
|
lambda y: get_column_str(
|
|
table_name=table_name, column_name=y[1]),
|
|
filter(
|
|
lambda y: y[0] == table_id,
|
|
zip(
|
|
db_column_names['table_id'],
|
|
db_column_names['column_name'],
|
|
),
|
|
),
|
|
)),
|
|
) for table_id, table_name in enumerate(db_table_names)
|
|
]
|
|
if schema_serialization_randomized:
|
|
random.shuffle(tables)
|
|
if schema_serialization_with_db_id:
|
|
serialized_schema = db_id_str.format(
|
|
db_id=db_id) + table_sep.join(tables)
|
|
else:
|
|
serialized_schema = table_sep.join(tables)
|
|
return serialized_schema
|
|
|
|
|
|
def form_input_for_construction(query, question, db_id, db_path, schema):
|
|
return {
|
|
'query':
|
|
query,
|
|
'question':
|
|
question,
|
|
'db_id':
|
|
db_id,
|
|
'db_path':
|
|
db_path,
|
|
'db_table_names':
|
|
schema['table_names_original'],
|
|
'db_column_names': {
|
|
'table_id': [
|
|
table_id
|
|
for table_id, column_name in schema['column_names_original']
|
|
],
|
|
'column_name': [
|
|
column_name
|
|
for table_id, column_name in schema['column_names_original']
|
|
]
|
|
},
|
|
'db_column_types':
|
|
schema['column_types'],
|
|
'db_primary_keys': [{
|
|
'column_id': column_id
|
|
} for column_id in schema['primary_keys']],
|
|
'db_foreign_keys': {
|
|
'column_id': [
|
|
column_id
|
|
for column_id, other_column_id in schema['foreign_keys']
|
|
],
|
|
'other_column_id': [
|
|
other_column_id
|
|
for column_id, other_column_id in schema['foreign_keys']
|
|
]
|
|
},
|
|
}
|