explore_banking77.py

# Created by xunannancy at 2024/03/04
import os
from dataset_utilities import dataset_folder_dict, label_counts_dict, prefix_dict, label_check
import json
from tqdm import tqdm
import random
import numpy as np
import csv
import pickle
import itertools
import subprocess
from copy import deepcopy
import gzip
import scipy.sparse as sp

def save_labels(dataset = 'BANKING77'):
    # label_title_id_mapping = pickle.load(open(f'{dataset_folder_dict[dataset]}/dpr_processed/labels/label_title_id_mapping.pkl','rb'))

    # label_list = list()
    # for line in open(f'{dataset_folder_dict[dataset]}/label_exploration/definitions_77_5examples_truncated.json'):
    #     label_list.append(json.loads(line)['title'])
    with open(f'./repositories/Few-Shot-Intent-Detection/Datasets/{dataset}/train_5/label') as f:
        label_list = f.readlines()
    label_list = sorted(list(set([i.strip().replace('_', ' ') for i in label_list])))
    os.makedirs(dataset_folder_dict[dataset], exist_ok=True)
    label_saved_path = f'{dataset_folder_dict[dataset]}/lbl.json.gz'
    # save lbl.json.gz
    fout = gzip.open(label_saved_path, 'w')
    for idx, label in enumerate(label_list):
        cur_str = json.dumps(
                    {
                        'uid': idx, #label_title_id_mapping[label],
                        'title': label,
                        'content': '',
                        'target_ind': [],
                        'target_rel': [],
                    }
                )+'\n'
        fout.write(cur_str.encode('utf-8'))
    fout.close()
    return

def retrieve_labels_sbert_train_examples(dataset, setting, batch_examples, model_name, num_examples=-1):
    # setting is the source of label
    assert setting in ['train_5', 'train_10']

    if dataset == 'BANKING77':
        mode = 'argmax'
    else:
        mode = 'average' # according to dev performance
    if batch_examples == 5:
        saved_path = os.path.join(dataset_folder_dict[dataset], 'label_exploration', f'{setting}_{model_name}_{num_examples}examples.json')
    else:
        saved_path = os.path.join(dataset_folder_dict[dataset], 'label_exploration', f'{setting}_{model_name}_{num_examples}examples_B{batch_examples}.json')
    os.makedirs(os.path.dirname(saved_path), exist_ok=True)

    with open(f'./repositories/Few-Shot-Intent-Detection/Datasets/{dataset}/train_5/label') as f:
        labels = f.readlines()
    labels = sorted(set([i.strip().replace('_', ' ') for i in labels]))
    assert len(labels) == label_counts_dict[dataset]

    with open(f'./repositories/Few-Shot-Intent-Detection/Datasets/{dataset}/train/seq.in') as f:
        seq = f.readlines()
    # NOTE: output from pred_sim in dpr_sim.py
    pred_path = f'./results/sbert_{mode}_{model_name}/{dataset}_train/{setting}_-1examples/trn/predictions_text.json'
    assert len(seq) == sum(1 for i in open(pred_path))
    label_instance_dict = dict()
    for cur_seq, pred_line in zip(seq, open(pred_path)):
        cur_label = json.loads(pred_line)[0].replace('_', ' ')
        if cur_label in label_instance_dict:
            label_instance_dict[cur_label].append(cur_seq)
        else:
            label_instance_dict[cur_label] = [cur_seq]
    assert len(label_instance_dict) == len(labels) # ensure every label has instances
    fout = open(saved_path, 'w')
    idx = 0
    for label in tqdm(labels, total=len(labels)):
        example_list = label_instance_dict[label]
        random.shuffle(example_list)
        num_batches = max(1, len(example_list)//batch_examples)
        # num_batches = int(np.ceil(len(example_list)/batch_examples))
        for batch_idx in range(num_batches):
            saved_res = {
                'id': idx,
                'title': label,
                'definition': ' '.join(example_list[batch_idx*batch_examples:(batch_idx+1)*batch_examples]),
            }
            fout.write(json.dumps(saved_res)+'\n')
            idx += 1
    fout.close()
    return

def prepare_data_sparse_retrieval(num_words=100, with_text=True,
                                  source_list=['definitions_77_5examples_truncated'],
                                  nickname='combined_5examples', dataset='BANKING77'):
    title_id_mapping, id_title_mapping = dict(), dict()

    data_folder = dataset_folder_dict[dataset]
    label_nickname = 'labels'
    if not with_text:
        assert len(source_list) == 0
        assert nickname == 'combined_without_text'
        # source = 'definitions_77.json'
        source = ''
    if num_words == 100:
        saved_file = f'{data_folder}/dpr_processed/{label_nickname}/{nickname}/combined.jsonl'
    else:
        saved_file = f'{data_folder}/dpr_processed/{label_nickname}/{nickname}_w{num_words}/combined.jsonl'

    file_path = f'{data_folder}/dpr_processed/{label_nickname}/{nickname}_w{num_words}.tsv'

    os.makedirs(os.path.dirname(saved_file), exist_ok=True)
    fout = open(saved_file, 'w')
    os.makedirs(os.path.dirname(file_path), exist_ok=True)

    f_write = open(file_path, 'w', newline='')
    writer = csv.writer(f_write, delimiter='\t', lineterminator='\n')
    writer.writerow(['id', 'text', 'title', 'label'])

    global_id = 1
    definitions = dict()
    title_set = set([])
    for source in source_list:
        definitions[source] = dict()
        num_lines = sum(1 for i in open(f'{data_folder}/label_exploration/{source}.json'))
        for line in tqdm(open(f'{data_folder}/label_exploration/{source}.json'), total=num_lines):
            row = json.loads(line)
            if row['title'] in definitions[source]:
                definitions[source][row['title']].append(row['definition'])
            else:
                definitions[source][row['title']] = [row['definition']]
            title_set.update([row['title']])

    if source == '':
        assert not with_text
        label_title_id_mapping = pickle.load(open(f'{dataset_folder_dict[dataset]}/dpr_processed/labels/label_title_id_mapping.pkl', 'rb'))
        for title in label_title_id_mapping:
            fout.write(
                json.dumps({
                    'id': global_id,
                    'contents': '',
                    'label': title,
                    'title': title,
                    'source': 'gpt-3.5-turbo',
                })+'\n')
            id_title_mapping[global_id] = [title]
            if title in title_id_mapping:
                title_id_mapping[title].append(global_id)
            else:
                title_id_mapping[title] = [global_id]

            writer.writerow([global_id, '', title, title]) # ['id', 'text', 'title', 'label']
            global_id += 1
    else:
        for title in tqdm(title_set, total=len(title_set)):
            if not with_text:
                fout.write(
                    json.dumps({
                        'id': global_id,
                        'contents': '',
                        'label': title,
                        'title': title,
                        'source': 'gpt-3.5-turbo',
                    })+'\n')
                id_title_mapping[global_id] = [title]
                if title in title_id_mapping:
                    title_id_mapping[title].append(global_id)
                else:
                    title_id_mapping[title] = [global_id]

                writer.writerow([global_id, '', title, title]) # ['id', 'text', 'title', 'label']
                global_id += 1
            else:
                all_text = ' '.join(list(itertools.chain.from_iterable([definitions[source][title] for source in source_list])))
                tokens = all_text.split()
                for chunk_id in range(max(1, len(tokens) // num_words)):
                    text = ' '.join(tokens[chunk_id * num_words:(chunk_id + 1) * num_words])
                    fout.write(
                        json.dumps({
                            'id': global_id,
                            'contents': text,
                            'label': title,
                            'title': title,
                            'source': 'gpt-3.5-turbo',
                        })+'\n')
                    id_title_mapping[global_id] = [title]
                    if title in title_id_mapping:
                        title_id_mapping[title].append(global_id)
                    else:
                        title_id_mapping[title] = [global_id]

                    writer.writerow([global_id, text.replace('\n', ' '), title, title]) # ['id', 'text', 'title', 'label']
                    global_id += 1

    # assert ids in id_title_mapping and title_id_mapping are the same set
    ids_ids, title_ids = np.sort(list(id_title_mapping.keys())), np.sort(list(itertools.chain(*list(title_id_mapping.values()))))
    assert np.array_equal(ids_ids, title_ids)

    print(f'global_id: {global_id}')
    print(f'title_id_mapping: {len(title_id_mapping)}')
    fout.close()
    f_write.close()

    if num_words == 100:
        with open(f'{data_folder}/dpr_processed/{label_nickname}/{nickname}_title_id_mapping.pkl', 'wb') as f:
            pickle.dump(title_id_mapping, f)
        with open(f'{data_folder}/dpr_processed/{label_nickname}/{nickname}_id_title_mapping.pkl', 'wb') as f:
            pickle.dump(id_title_mapping, f)
    else:
        with open(f'{data_folder}/dpr_processed/{label_nickname}/{nickname}_w{num_words}_title_id_mapping.pkl', 'wb') as f:
            pickle.dump(title_id_mapping, f)
        with open(f'{data_folder}/dpr_processed/{label_nickname}/{nickname}_w{num_words}_id_title_mapping.pkl', 'wb') as f:
            pickle.dump(id_title_mapping, f)
    return

def construct_dataset(dataset, label_nickname='combined_definition_5examples', max_context_token=256, parent_dataset='BANKING77', seed=0,
                      with_question_flag=False):
    random.seed(seed)
    if with_question_flag:
        prefix = prefix_dict[parent_dataset]
    else:
        prefix = ''

    with open(f'./repositories/Few-Shot-Intent-Detection/Datasets/{parent_dataset}/{dataset}/seq.in') as f:
        queries = f.readlines()
    queries = [i.strip() for i in queries]
    with open(f'./repositories/Few-Shot-Intent-Detection/Datasets/{parent_dataset}/{dataset}/label') as f:
        labels = f.readlines()
    labels = [i.strip().replace('_', ' ') for i in labels]

    combined_title_id_mapping = pickle.load(open(f'{dataset_folder_dict[parent_dataset]}/dpr_processed/labels/{label_nickname}_title_id_mapping.pkl', 'rb'))
    combined_doc = dict()
    for line in tqdm(open(f'{dataset_folder_dict[parent_dataset]}/dpr_processed/labels/{label_nickname}/combined.jsonl')):
        info = json.loads(line)
        id, contents = info['id'], info['contents']
        combined_doc[int(id)] = deepcopy(contents)
    label_title_id_mapping = pickle.load(open(f'{dataset_folder_dict[parent_dataset]}/dpr_processed/labels/label_title_id_mapping.pkl', 'rb'))

    saved_instances = list()
    assert set(labels) == set(list(combined_title_id_mapping.keys()))
    for query, label in zip(queries, labels):
        positive_ctxs = list()
        for doc_id in combined_title_id_mapping[label.replace('_', ' ')]:
            text = combined_doc[doc_id]
            if query in text:
                # NOTE: mainly for combined_train_5_5examples, combined_train_5examples, combined_train_10_5examples
                text = text[:text.find(query)]+text[text.find(query)+len(query):]
            positive_ctxs.append({
                'title': label,
                'text': text,
                'passage_id': doc_id,
            })

        cur_instance = {
            'question': prefix + query,
            'answers': [label],
            'positive_ctxs': positive_ctxs,
            'negative_ctxs': [],
            'hard_negatives_ctxs': [],
        }
        saved_instances.append(cur_instance)

    simple_name_mapping = {
        'valid': 'dev',
        'test': 'tst',
        'train_5': 'trn',
        'train_10': 'trn',
        'train': 'trn',
    }

    suffix = label_nickname[len('combined'):]
    if with_question_flag:
        suffix += '_question'
    if dataset in ['valid', 'test']:
        folder_name = 'train_5'
    elif dataset in ['train_5', 'train_10', 'train']:
        folder_name = dataset
    saved_folder = f'{dataset_folder_dict[f"{parent_dataset}_{folder_name}"]}/dpr_processed/query_context_direct_{max_context_token}{suffix}'
    os.makedirs(saved_folder, exist_ok=True)
    saved_path = f'{saved_folder}/{simple_name_mapping[dataset]}.json'
    with open(saved_path, 'w') as f:
        json.dump(saved_instances, f, indent=4)
    fout = gzip.open(f'{dataset_folder_dict[parent_dataset]}/{folder_name}/{simple_name_mapping[dataset]}.json.gz', 'w')
    y_cols, y_rows = list(), list()
    num_instances = 0
    for idx, instance in enumerate(saved_instances):
        cur_labels = list(itertools.chain(*[label_title_id_mapping[i] for i in instance['answers']]))
        output = json.dumps({
            'uid': idx,
            'title': str(idx),
            'content': instance['question'],
            'target_ind': cur_labels,
        })+'\n'
        fout.write(output.encode('utf-8'))
        for label in cur_labels:
            y_cols.append(label)
            y_rows.append(idx)
        num_instances += 1
    fout.close()

    saved_y_path = f'{dataset_folder_dict[parent_dataset]}/{folder_name}/processed/Y.{simple_name_mapping[dataset]}.npz'
    os.makedirs(os.path.dirname(saved_y_path), exist_ok=True)
    cur_y = sp.coo_matrix(
        (np.ones_like(y_cols), (y_rows, y_cols)),
        shape=(num_instances, len(label_title_id_mapping)),
    )
    sp.save_npz(saved_y_path, cur_y.tocsr())
    if dataset in ['valid', 'test']:
        for available_folder_name in ['train_10', 'train']:
            if os.path.exists(f'{dataset_folder_dict[parent_dataset]}/{available_folder_name}/{simple_name_mapping[dataset]}.json.gz'):
                os.remove(f'{dataset_folder_dict[parent_dataset]}/{available_folder_name}/{simple_name_mapping[dataset]}.json.gz')
            des_file = f'{dataset_folder_dict[parent_dataset]}/{available_folder_name}/{simple_name_mapping[dataset]}.json.gz'
            os.makedirs(os.path.dirname(des_file), exist_ok=True)
            subprocess.run(['ln', '-s', f'{dataset_folder_dict[parent_dataset]}/train_5/{simple_name_mapping[dataset]}.json.gz', des_file])
            if os.path.exists(f'{dataset_folder_dict[parent_dataset]}/{available_folder_name}/processed/Y.{simple_name_mapping[dataset]}.npz'):
                os.remove(f'{dataset_folder_dict[parent_dataset]}/{available_folder_name}/processed/Y.{simple_name_mapping[dataset]}.npz')
            os.makedirs(os.path.dirname(saved_y_path.replace('train_5', available_folder_name)), exist_ok=True)
            des_file = saved_y_path.replace('train_5', available_folder_name)
            os.makedirs(os.path.dirname(des_file), exist_ok=True)
            subprocess.run(['ln', '-s', saved_y_path, des_file])
            if os.path.exists(saved_path.replace('train_5', available_folder_name)):
                os.remove(saved_path.replace('train_5', available_folder_name))
            des_file = saved_path.replace('train_5', available_folder_name)
            os.makedirs(os.path.dirname(des_file), exist_ok=True)
            subprocess.run(["ln", "-s", saved_path, des_file])
    else:
        saved_folder = f'{dataset_folder_dict[f"{parent_dataset}_{folder_name}"]}/dpr_processed/query_context_direct_{max_context_token}{suffix}'
        os.makedirs(saved_folder, exist_ok=True)

        if dataset in ['train', 'train_5', 'train_10']:
            label_instances = dict()
            for idx, instance in enumerate(saved_instances):
                if instance['answers'][0] not in label_instances:
                    label_instances[instance['answers'][0]] = [idx]
                else:
                    label_instances[instance['answers'][0]].append(idx)
            if dataset in ['train_5', 'train_10']:
                # each label has fixed number of instances
                instance_indices = np.array(list(label_instances.values())).transpose().reshape([-1])
                saved_instances = [saved_instances[i] for i in instance_indices]
        else:
            random.shuffle(saved_instances)
        with open(f'{saved_folder}/trn.json', 'w') as f:
            json.dump(saved_instances, f, indent=4)

    return

if __name__ == '__main__':
    save_labels()
    label_check(dataset='BANKING77')

    model_name = 'gtr-t5-xl'
    for setting in ['train_5', 'train_10']:
        retrieve_labels_sbert_train_examples(dataset='BANKING77', setting=setting, batch_examples=5, model_name=model_name)
        for num_words in [100]:
            prepare_data_sparse_retrieval(num_words=num_words, with_text=True,
                                          source_list=[f'{setting}_{model_name}_-1examples'],
                                          nickname=f'combined_{setting}_{model_name}_-1examples')

    for dataset in ['train_5', 'test', 'valid', 'train_10', 'train']:
        for label_nickname in [f'combined_train_5_{model_name}_-1examples', f'combined_train_10_{model_name}_-1examples']:
            print(f'{dataset}-{label_nickname}...')
            for with_question_flag in [False]: #[True, False]:
                construct_dataset(dataset=dataset, label_nickname=label_nickname, with_question_flag=with_question_flag)