<filename>data_synthesis/augment_kra.py
from random import shuffle
from typing import List, Dict, Any


if __name__ == '__main__':
    with open('/users10/dzrwang/SCoRE/score/data/augment_all_context.txt', 'r', encoding='utf-8') as f:
        data: List[Dict[str, Any]] = []
        for line in f:
            line = line.strip().split(' ||| ')
            query_schema: List[str] = line[0].split(' </s> ')

            one_data: Dict[str, Any] = {'query': query_schema[0],
                                        'schema': ' </s> '.join(query_schema[1:]),
                                        'column_type': line[1]}
            one_data['column_type'] = [
                int(x) for x in one_data['column_type'].split()]
            one_data['kra'] = [0 for _ in range(len(one_data['column_type']))]
            if len(line) == 3:
                one_data['action'] = line[2]

            data.append(one_data)

    data.sort(key=lambda x: x['query'][::-1])
    for i in range(1, len(data)):
        last_data = data[i-1]
        current_data = data[i]

        if not(current_data['query'].endswith(last_data['query'])) or last_data['query'] == current_data['query']:
            for j in range(len(current_data['kra'])):
                if current_data['column_type'][j] > 0:
                    current_data['kra'][j] = 2
            data[i] = current_data
            continue
        assert len(last_data['kra']) == len(current_data['kra']), print(
            len(last_data['kra']), len(current_data['kra']))

        for j in range(len(last_data['kra'])):
            if last_data['column_type'][j] > 0:
                if current_data['column_type'][j] == 0:
                    current_data['kra'][j] = 1
            elif last_data['column_type'][j] == 0:
                if current_data['column_type'][j] > 0:
                    current_data['kra'][j] = 2

        data[i] = current_data
    
    for i in range(len(data)):
        data[i]['kra'] = [str(x) for x in data[i]['kra']]
        data[i]['column_type'] = [str(x) for x in data[i]['column_type']]
        if 'action' in data[i]:
            data[i]['action'] = [str(x) for x in data[i]['action']]

    shuffle(data)
    with open('/users10/dzrwang/SCoRE/score/data/augment_all_context_with_kra.txt', 'w', encoding='utf-8') as f:
        for d in data:
            if 'action' in data:
                current_data = f'{d["query"]}  </s> {d["schema"]} ||| {" ".join(d["column_type"])} ||| {" ".join(d["action"])} ||| {" ".join(d["kra"])}'
            else:
                current_data = f'{d["query"]}  </s> {d["schema"]} ||| {" ".join(d["column_type"])} ||| {" ".join(d["kra"])}'
            f.write(current_data + '\n')

    print('augment done')

<filename>data_synthesis/augment_kra.py
from random import shuffle
from typing import List, Dict, Any


if __name__ == '__main__':
    with open('/users10/dzrwang/SCoRE/score/data/augment_all_context.txt', 'r', encoding='utf-8') as f:
        data: List[Dict[str, Any]] = []
        for line in f:
            line = line.strip().split(' ||| ')
            query_schema: List[str] = line[0].split(' </s> ')

            one_data: Dict[str, Any] = {'query': query_schema[0],
                                        'schema': ' </s> '.join(query_schema[1:]),
                                        'column_type': line[1]}
            one_data['column_type'] = [
                int(x) for x in one_data['column_type'].split()]
            one_data['kra'] = [0 for _ in range(len(one_data['column_type']))]
            if len(line) == 3:
                one_data['action'] = line[2]

            data.append(one_data)

    data.sort(key=lambda x: x['query'][::-1])
    for i in range(1, len(data)):
        last_data = data[i-1]
        current_data = data[i]

        if not(current_data['query'].endswith(last_data['query'])) or last_data['query'] == current_data['query']:
            for j in range(len(current_data['kra'])):
                if current_data['column_type'][j] > 0:
                    current_data['kra'][j] = 2
            data[i] = current_data
            continue
        assert len(last_data['kra']) == len(current_data['kra']), print(
            len(last_data['kra']), len(current_data['kra']))

        for j in range(len(last_data['kra'])):
            if last_data['column_type'][j] > 0:
                if current_data['column_type'][j] == 0:
                    current_data['kra'][j] = 1
            elif last_data['column_type'][j] == 0:
                if current_data['column_type'][j] > 0:
                    current_data['kra'][j] = 2

        data[i] = current_data
    
    for i in range(len(data)):
        data[i]['kra'] = [str(x) for x in data[i]['kra']]
        data[i]['column_type'] = [str(x) for x in data[i]['column_type']]
        if 'action' in data[i]:
            data[i]['action'] = [str(x) for x in data[i]['action']]

    shuffle(data)
    with open('/users10/dzrwang/SCoRE/score/data/augment_all_context_with_kra.txt', 'w', encoding='utf-8') as f:
        for d in data:
            if 'action' in data:
                current_data = f'{d["query"]}  </s> {d["schema"]} ||| {" ".join(d["column_type"])} ||| {" ".join(d["action"])} ||| {" ".join(d["kra"])}'
            else:
                current_data = f'{d["query"]}  </s> {d["schema"]} ||| {" ".join(d["column_type"])} ||| {" ".join(d["kra"])}'
            f.write(current_data + '\n')

    print('augment done')
