<reponame>WDZRMPCBIT/chase<filename>IGSQL/split_output.py
import json
from functools import cmp_to_key
from typing import Dict, List, Any


if __name__ == '__main__':
    with open('./output/mix/valid_use_predicted_queries_predictions.json', 'r', encoding='utf-8') as f:
        data: List[Dict[str, Any]] = []
        for line in f:
            data.append(json.loads(line))
    with open('./output/mix/map.json', 'r', encoding='utf=8') as f:
        database_map: Dict[str, Dict[int, str]] = json.load(f)

    better_database: Dict[str, List[Dict[str, Any]]] = {
        'spider': [], 'cosql': [], 'sparc': []}
    for d in data:
        better_database[database_map[d['database_id']]
                        [d['interaction_id']]].append(d)
    for name in better_database:
        better_database[name].sort(key=lambda x: (x['database_id'], int(x['interaction_id'])))

    database: Dict[str, List[Dict[str, Any]]] = {
        'spider': [], 'cosql': [], 'sparc': []}
    for name in database:
        with open(f'./output/{name}/valid_use_predicted_queries_predictions.json', 'r', encoding='utf-8') as f:
            data = []
            for line in f:
                data.append(json.loads(line))
        data.sort(key=lambda x: (x['database_id'], int(x['interaction_id'])))
        database[name] = data

    for name in better_database:
        assert len(database[name]) == len(better_database[name]), print(
            name, len(database[name]), len(better_database[name]))
        for i in range(len(database[name])):
            assert (len(database[name][i]['input_seq']) == len(better_database[name][i]['input_seq'])), print(
                str(database[name][i]) + '\n' + str(better_database[name][i]))
            database[name][i]['probability'] = better_database[name][i]['probability']
            database[name][i]['prediction'] = better_database[name][i]['prediction']
            database[name][i]['flat_prediction'] = better_database[name][i]['flat_prediction']

        with open(f'./output/{name}/valid_use_predicted_queries_predictions.json', 'w', encoding='utf-8') as f:
            for d in database[name]:
                f.write(json.dumps(d) + '\n')

    print('split done')

import json
from functools import cmp_to_key
from typing import Dict, List, Any


if __name__ == '__main__':
    with open('./output/mix/valid_use_predicted_queries_predictions.json', 'r', encoding='utf-8') as f:
        data: List[Dict[str, Any]] = []
        for line in f:
            data.append(json.loads(line))
    with open('./output/mix/map.json', 'r', encoding='utf=8') as f:
        database_map: Dict[str, Dict[int, str]] = json.load(f)

    better_database: Dict[str, List[Dict[str, Any]]] = {
        'spider': [], 'cosql': [], 'sparc': []}
    for d in data:
        better_database[database_map[d['database_id']]
                        [d['interaction_id']]].append(d)
    for name in better_database:
        better_database[name].sort(key=lambda x: (x['database_id'], int(x['interaction_id'])))

    database: Dict[str, List[Dict[str, Any]]] = {
        'spider': [], 'cosql': [], 'sparc': []}
    for name in database:
        with open(f'./output/{name}/valid_use_predicted_queries_predictions.json', 'r', encoding='utf-8') as f:
            data = []
            for line in f:
                data.append(json.loads(line))
        data.sort(key=lambda x: (x['database_id'], int(x['interaction_id'])))
        database[name] = data

    for name in better_database:
        assert len(database[name]) == len(better_database[name]), print(
            name, len(database[name]), len(better_database[name]))
        for i in range(len(database[name])):
            assert (len(database[name][i]['input_seq']) == len(better_database[name][i]['input_seq'])), print(
                str(database[name][i]) + '\n' + str(better_database[name][i]))
            database[name][i]['probability'] = better_database[name][i]['probability']
            database[name][i]['prediction'] = better_database[name][i]['prediction']
            database[name][i]['flat_prediction'] = better_database[name][i]['flat_prediction']

        with open(f'./output/{name}/valid_use_predicted_queries_predictions.json', 'w', encoding='utf-8') as f:
            for d in database[name]:
                f.write(json.dumps(d) + '\n')

    print('split done')