md_bayesian_er_ditto/ml_er/ditto_er.py

import pickle
import torch
import json
import numpy as np
import random
# from ditto.matcher import *
from setting import *
from colorama import Fore
from argparse import Namespace
import ConfigSpace
from ConfigSpace import Configuration
from ditto.matcher import set_seed, to_str, classify, predict, tune_threshold, load_model
from ConfigSpace.read_and_write import json as csj
from ditto.ditto_light.dataset import DittoDataset
from ditto.ditto_light.summarize import Summarizer
from ditto.ditto_light.knowledge import *
from ditto.ditto_light.ditto import train
import os
import sys

os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
sys.path.append('/root/hjt/md_bayesian_er_ditto/')


def matching(config):
    print(Fore.BLUE + f'Config: {config}')

    # with open(md_output_dir + r"\mds.pickle", "rb") as file:
    #     md_list = pickle.load(file)

    hp = Namespace()
    hp.task = directory_path.replace('/root/hjt/DeepMatcher Dataset/', '')
    hp.run_id = 0
    hp.batch_size = config['batch_size']
    hp.max_len = config['max_len']
    hp.lr = 3e-5
    hp.n_epochs = 20
    # hp.finetuning
    hp.save_model = True
    hp.input_path = '/root/autodl-tmp/input/candidates_small.jsonl'
    hp.output_path = '/root/autodl-tmp/output/matched_small.jsonl'
    hp.logdir = '/root/autodl-tmp/checkpoints/'
    hp.checkpoint_path = '/root/autodl-tmp/checkpoints/'
    hp.lm = config['language_model']
    hp.fp16 = config['half_precision_float']
    hp.da = config['data_augmentation']
    hp.alpha_aug = 0.8
    hp.dk = None
    hp.summarize = config['summarize']
    hp.size = None
    hp.use_gpu = True

    seed = hp.run_id
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

    # only a single task for baseline
    task = hp.task

    # create the tag of the run
    run_tag = '%s_lm=%s_da=%s_dk=%s_su=%s_size=%s_id=%d' % (task, hp.lm, hp.da,
                                                            hp.dk, hp.summarize, str(hp.size), hp.run_id)
    run_tag = run_tag.replace('/', '_')

    # load task configuration
    configs = json.load(open('configs.json'))
    configs = {conf['name']: conf for conf in configs}
    config = configs[task]

    trainset = config['trainset']
    validset = config['validset']
    testset = config['testset']

    # summarize the sequences up to the max sequence length
    if hp.summarize:
        summarizer = Summarizer(config, lm=hp.lm)
        trainset = summarizer.transform_file(trainset, max_len=hp.max_len)
        validset = summarizer.transform_file(validset, max_len=hp.max_len)
        testset = summarizer.transform_file(testset, max_len=hp.max_len)

    # load train/dev/test sets
    train_dataset = DittoDataset(trainset,
                                 lm=hp.lm,
                                 max_len=hp.max_len,
                                 size=hp.size,
                                 da=hp.da)
    valid_dataset = DittoDataset(validset, lm=hp.lm)
    test_dataset = DittoDataset(testset, lm=hp.lm)

    # train and evaluate the model
    train(train_dataset,
          valid_dataset,
          test_dataset,
          run_tag, hp)

    set_seed(123)
    config, model = load_model(hp.task, hp.checkpoint_path,
                               hp.lm, hp.use_gpu, hp.fp16)

    summarizer = dk_injector = None
    if hp.summarize:
        summarizer = Summarizer(config, hp.lm)

    # tune threshold
    threshold = tune_threshold(config, model, hp)

    # run prediction
    predict(hp.input_path, hp.output_path, config, model,
            summarizer=summarizer,
            max_len=hp.max_len,
            lm=hp.lm,
            dk_injector=dk_injector,
            threshold=threshold)
    # todo indicators
    #  write results
    #  interpretability


# todo ml_er function
def ml_er(config: Configuration):
    indicators = matching(config)
    output_path = er_output_dir + r"\eval_result.txt"
    with open(output_path, 'w') as _f:
        _f.write('F1:' + str(indicators["F1"]) + '\n')
        _f.write('interpretability:' + str(indicators['interpretability']) + '\n')
        _f.write('performance:' + str(indicators['performance']) + '\n')


if __name__ == '__main__':
    if os.path.isfile(hpo_output_dir + r"\incumbent.json"):
        with open(hpo_output_dir + r"\configspace.json", 'r') as f:
            dict_configspace = json.load(f)
        str_configspace = json.dumps(dict_configspace)
        configspace = csj.read(str_configspace)
        with open(hpo_output_dir + r"\incumbent.json", 'r') as f:
            dic = json.load(f)
        configuration = ConfigSpace.Configuration(configspace, values=dic)
        ml_er(configuration)