| """Training data prep utils.""" | |
| import json | |
| import re | |
| from collections import defaultdict | |
| from schema import ForeignKey, Table, TableColumn | |
| def read_tables_json( | |
| schema_file: str, | |
| lowercase: bool = False, | |
| ) -> dict[str, dict[str, Table]]: | |
| """Read tables json.""" | |
| data = json.load(open(schema_file)) | |
| db_to_tables = {} | |
| for db in data: | |
| db_name = db["db_id"] | |
| table_names = db["table_names_original"] | |
| db["column_names_original"] = [ | |
| [x[0], x[1]] for x in db["column_names_original"] | |
| ] | |
| db["column_types"] = db["column_types"] | |
| if lowercase: | |
| table_names = [tn.lower() for tn in table_names] | |
| pks = db["primary_keys"] | |
| fks = db["foreign_keys"] | |
| tables = defaultdict(list) | |
| tables_pks = defaultdict(list) | |
| tables_fks = defaultdict(list) | |
| for idx, ((ti, col_name), col_type) in enumerate( | |
| zip(db["column_names_original"], db["column_types"]) | |
| ): | |
| if ti == -1: | |
| continue | |
| if lowercase: | |
| col_name = col_name.lower() | |
| col_type = col_type.lower() | |
| if idx in pks: | |
| tables_pks[table_names[ti]].append( | |
| TableColumn(name=col_name, dtype=col_type) | |
| ) | |
| for fk in fks: | |
| if idx == fk[0]: | |
| other_column = db["column_names_original"][fk[1]] | |
| other_column_type = db["column_types"][fk[1]] | |
| other_table = table_names[other_column[0]] | |
| tables_fks[table_names[ti]].append( | |
| ForeignKey( | |
| column=TableColumn(name=col_name, dtype=col_type), | |
| references_name=other_table, | |
| references_column=TableColumn( | |
| name=other_column[1], dtype=other_column_type | |
| ), | |
| ) | |
| ) | |
| tables[table_names[ti]].append(TableColumn(name=col_name, dtype=col_type)) | |
| db_to_tables[db_name] = { | |
| table_name: Table( | |
| name=table_name, | |
| columns=tables[table_name], | |
| pks=tables_pks[table_name], | |
| fks=tables_fks[table_name], | |
| examples=None, | |
| ) | |
| for table_name in tables | |
| } | |
| return db_to_tables | |
| def clean_str(target: str) -> str: | |
| """Clean string for question.""" | |
| if not target: | |
| return target | |
| target = re.sub(r"[^\x00-\x7f]", r" ", target) | |
| line = re.sub(r"''", r" ", target) | |
| line = re.sub(r"``", r" ", line) | |
| line = re.sub(r"\"", r"'", line) | |
| line = re.sub(r"[\t ]+", " ", line) | |
| return line.strip() | |