Function generate_db_prompt_bird_v2

data_preprocess.py:461–506

:param root_place: :param db_name: :return:

(root_dir, dataset, db_id, limit_value=3)

def generate_db_prompt_bird_v2(root_dir, dataset, db_id, limit_value=3):
    # extract create ddls
    '''
    :param root_place:
    :param db_name:
    :return:
    '''
    DATA_PATH = root_dir
    dev_db_path = f"{DATA_PATH}/{dataset}/dev_databases"
    db_uri = dev_db_path + "/" + db_id + "/" + db_id + ".sqlite"
    full_schema_prompt_list = []
    table_names = []
    conn = sqlite3.connect(db_uri)
    # Create a cursor object
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = cursor.fetchall()
    schemas = {}
    for table in tables:
        if table == 'sqlite_sequence':
            continue
        if isinstance(table, tuple):
            table_name = table[0]
            table_names.append(table_name)
        cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name='{}';".format(table[0]))
        create_prompt = cursor.fetchone()[0]
        schemas[table[0]] = create_prompt
        if limit_value:
            cur_table = table[0]
            if cur_table in ['order', 'by', 'group']:
                cur_table = "`{}`".format(cur_table)

            cursor.execute("SELECT * FROM {} LIMIT {}".format(cur_table, limit_value))
            column_names = [description[0] for description in cursor.description]
            values = cursor.fetchall()
            rows_prompt = nice_look_table(column_names=column_names, values=values)
            verbose_prompt = "/* \n {} example rows: \n SELECT * FROM {} LIMIT {}; \n {} \n */".format(
                limit_value, cur_table, limit_value, rows_prompt)
            schemas[table[0]] = "{} \n {}".format(create_prompt, verbose_prompt)

    for k, v in schemas.items():
        full_schema_prompt_list.append(v)

    schema_prompt = "\n\n".join(full_schema_prompt_list)

    return schema_prompt, table_names


if __name__ == "__main__":