demo.py

# coding=utf-8
# can be used for demo or test
import os
import sys
import json5
import codecs
from src.utils import params, loader
from src.demoer import Demoer
from src.utils.metrics import registry as metrics
from curLine_file import curLine, normal_transformer

def main():
    argv = sys.argv
    host_name = argv[2]
    model_id = argv[3]
    print(curLine(), "argv:", argv)
    arg_groups = params.parse(argv[1], host_name, mode="test")
    args, config = arg_groups[0]
    args.output_dir = "/home/%s/Mywork/model/qa_model_dir/on_test/block1-layer1-hidden100-acc=85.31" % (host_name)  # TODO
    args.output_dir = "/home/%s/Mywork/model/qa_model_dir/part_chatcorpus_model/block1-layer1-hidden100-normal-acc80.57" % host_name

    args.data_dir = os.path.join("/home/%s/Mywork/corpus/Chinese_QA" % host_name, args.data_dir)
    checkpoint_dir = os.path.join(args.output_dir, model_id)

    if len(argv) == 5:
        args.eval_file = argv[4]
    demoer = Demoer(args, checkpoint_dir)
    sample = {'text1':"请问谁有狂三这张高清的电影资源？", 'text2': '这张高清图，谁有狂三这张高清的请问谁有狂三这张高清的电影资源？'}
    predictions, probabilities, inference_time = demoer.serve(dev=[sample])

    test(args, config, demoer) # 批量测试

    # predict_batchsize = 30
    # infer_flag = True # False  #
    #
    # if infer_flag:
    #     text2_list = []
    #     for i in range(predict_batchsize):
    #         text2_list.append(sample['text2'])
    #     batch = [{'text1': sample['text1'], 'text2_list': text2_list}]
    # else:
    #     batch = []
    #     for i in range(predict_batchsize):
    #         batch.append({'text1': sample['text1'], 'text2': sample['text2']})
    #
    # inference_time_sum1 = 0
    # for i in range(20):
    #     predictions, probabilities, inference_time = demoer.serve(dev=batch, infer_flag=infer_flag)
    #     inference_time_sum1 += inference_time
    # inference_time_sum =0
    # cishu = 400
    # for i in range(cishu):
    #     predictions, probabilities, inference_time = demoer.serve(dev=batch, infer_flag=infer_flag)
    #     inference_time_sum += inference_time
    #     # print(curLine(), inference_time)
    #     # print(curLine(), "predictions:",predictions)
    #     # print(curLine(), "probabilities:", probabilities[0])
    # print(curLine(), "inference_time1=%f ms, inference_time=%f ms" % (inference_time_sum1/20, inference_time_sum/cishu ))
    # print(curLine(), f"infer_flag:{infer_flag}, predict_batchsize={predict_batchsize}, inference_time={inference_time_sum/cishu}ms")

def test(args, config, demoer):
    dev = loader.load_data(args.data_dir, args.eval_file)
    targets = []
    for sample in dev:
        targets.append(int(sample['target']))
    predictions, probabilities, inference_time = demoer.serve(dev=dev, batch_size=384)

    if "train" in args.eval_file:  # 将模型的置信度保存到文件
        with open(os.path.join(args.data_dir, "%s.txt" % args.eval_file), "r") as fr:
            lines = fr.readlines()
        assert len(lines) == len(probabilities), 'number of lines is %d, number of probabilities is %d' % (len(lines), len(probabilities))
        save_file = os.path.join(args.data_dir, "%s_score.txt" % args.eval_file)
        with open(save_file, "w") as writer:
            for line, prediction, prob in zip(lines, predictions, probabilities):
                writer.write("%s\t%f\n" % (line.strip(), prob[1]))
        print(curLine(), "save %d results to %s" % (len(probabilities), save_file))

    outputs = {
        'target': targets,
        'prob': probabilities,
        'pred': predictions,
        'args': args,
    }
    # total_loss = sum(losses[:-1]) / (len(losses) - 1) if len(losses) > 1 else sum(losses)
    states = {
        'inference_time': inference_time/len(targets)
    }
    for metric in args.watch_metrics:
        if metric not in states:  # multiple metrics could be computed by the same function
            states.update(metrics[metric](outputs))
    print(curLine(), "stats:", states)
    with open('%s/log.jsonl'%args.output_dir, 'a') as f:
        f.write(json5.dumps({
            'data': os.path.basename(args.data_dir),
            'params': config,
            'state': states}))
        f.write('\n')

if __name__ == '__main__':
    main()