<filename>examples/bart/summarize.py<gh_stars>0
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch
from fairseq.models.bart import BARTModel
import argparse
import pickle
from multiprocessing import Pool
import os
import copy
import time
import pynvml

XSUM_KWARGS = dict(beam=6, lenpen=1.0, max_len_b=60, min_len=10, no_repeat_ngram_size=3)
CNN_KWARGS = dict(beam=4, lenpen=2.0, max_len_b=140, min_len=55, no_repeat_ngram_size=3)


@torch.no_grad()
def generate(bart, infile, outfile="bart_hypo.txt", ie_filename="test.source.binie",
                                        bsz=32, n_obs=None, **eval_kwargs):
    count = 1

    # if n_obs is not None: bsz = min(bsz, n_obs)
        
    with open(infile) as source, open(outfile, "w") as fout, open(ie_filename,"rb") as fie:
        s_ie = pickle.load(fie)

        sline = source.readline().strip()
        slines = [sline]
        source_ie = [s_ie[0]]
        ie_num_total = []
        for sline in source:
            if n_obs is not None and count > n_obs:
                break
            if count % bsz == 0:
                hypotheses_batch,extra = bart.sample(slines, source_ie, **eval_kwargs)
                ie_num_total.extend(extra)
                for hypothesis in hypotheses_batch:
                    fout.write(hypothesis + "\n")
                    fout.flush()
                slines = []

            slines.append(sline.strip())
            source_ie.append(s_ie[count])
            count += 1

        if slines != []:
            hypotheses_batch,extra = bart.sample(slines, source_ie,**eval_kwargs)
            for hypothesis in hypotheses_batch:
                fout.write(hypothesis + "\n")
                fout.flush()
    with open(outfile+"ienum","wb") as file:
        pickle.dump(ie_num_total,file)

def main(args):
    torch.cuda.set_device(args.gpu)
    eval_kwargs = XSUM_KWARGS if args.xsum_kwargs else CNN_KWARGS
    if args.model_dir == "pytorch/fairseq":
        bart = torch.hub.load("pytorch/fairseq", args.model_file)
    else:
        bart = BARTModel.from_pretrained(
            args.model_dir,
            checkpoint_file=args.model_file,
            data_name_or_path=args.model_dir,
        )
    bart = bart.eval()
    if torch.cuda.is_available():
        bart = bart.cuda().half()
    generate(
        bart, args.src, bsz=args.bsz, n_obs=args.n, outfile=args.out,ie_filename=args.src_ie, **eval_kwargs
    )

    return True

def multi_process(args,N_cut = 8):

    gpu_no = gpu_assign(N_cut)
    args_list = []

    split_dirtory = args.split_dirtory

    for i in range(len(gpu_no)):
        atmp = copy.deepcopy(args)
        atmp.gpu = gpu_no[i]
        atmp.src = os.path.join(split_dirtory,"test.source"+str(i))
        atmp.out = os.path.join(split_dirtory,"test.hypo"+str(i))
        atmp.src_ie = os.path.join(split_dirtory,"test.source.binie"+str(i))
        args_list.append(atmp)

    torch.multiprocessing.set_start_method("spawn",force=True)
    with Pool(len(gpu_no)) as pool:
        x_it = pool.imap_unordered(main, args_list)
        for i in x_it:
            print(i)
    mergy_data(split_dirtory,args.out,N_cut)
def split_data(args,N_cut = 8):
    
    infile = args.src
    outfile = args.out
    ie_filename=args.src_ie
    split_dirtory = args.split_dirtory
        
    with open(infile) as source, open(outfile, "w") as fout, open(ie_filename,"rb") as fie:
        
        s_ie = pickle.load(fie)
        sline = source.readlines()

        N_batch = int(len(sline)/N_cut)+2
        for i in range(N_cut):
            with open(os.path.join(split_dirtory,"test.source"+str(i)),"w") as file:
                file.write("".join(sline[i*N_batch:(i+1)*N_batch]))
            with open(os.path.join(split_dirtory,"test.source.binie"+str(i)),"wb") as file:
                pickle.dump(s_ie[i*N_batch:(i+1)*N_batch],file)

def mergy_data(in_dirtory,out_file,N_cut = 8):
    total_data = []
    total_ie = []

    for i in range(N_cut):

        with open(os.path.join(in_dirtory,"test.hypo"+str(i)),"r") as file:
            total_data.extend(file.readlines())

        with open(os.path.join(in_dirtory,"test.hypo"+str(i)+"ienum"),"rb") as file:
            total_ie.extend(pickle.load(file))

    with open(out_file,"w") as file:
        file.write("".join(total_data))

    with open(out_file+"ienum","wb") as file:
        pickle.dump(total_ie,file)


def gpu_assign(N_process):
    pynvml.nvmlInit()
    deviceCount = pynvml.nvmlDeviceGetCount()

    gpu_avi = []
    for i in range(deviceCount):
        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
        meminfo = pynvml.nvmlDeviceGetMemoryInfo(handle)
        mem_used = int(meminfo.used / (1024 ** 2) // 1024) + 1
        mem_free = int(meminfo.free / (1024 ** 2) // 1024)
        gpu_avi.append(mem_free)

    out = []
    for i in range(N_process):

        if max(gpu_avi) < 5:
            raise("Error")
        max_avi = gpu_avi.index(max(gpu_avi))
        out.append(max_avi)
        gpu_avi[max_avi] -= 5

    print("GPU SET INDEX",out)

    return out



if __name__ == "__main__":
    """
    Usage::

         python examples/bart/summarize.py \
            --model-dir $HOME/bart.large.cnn \
            --model-file model.pt \
            --src $HOME/data-bin/cnn_dm/test.source
    """
    start_time = time.time()
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model-dir",
        required=True,
        type=str,
        default="bart.large.cnn/",
        help="path containing model file and src_dict.txt",
    )
    parser.add_argument(
        "--model-file",
        default="checkpoint_best.pt",
        help="where in model_dir are weights saved",
    )
    parser.add_argument(
        "--src", default="test.source", help="text to summarize", type=str
    )
    parser.add_argument(
        "--out", default="test.hypo", help="where to save summaries", type=str
    )
    parser.add_argument(
        "--src_ie", default="test.source.binie", help="where to save summaries", type=str
    )
    parser.add_argument("--bsz", default=32, help="where to save summaries", type=int)
    parser.add_argument("--gpu", default=0, help="which gpu to use", type=int)
    parser.add_argument(
        "--n", default=None, help="how many examples to summarize", type=int
    )
    parser.add_argument(
        "--split_dirtory", default=".", help="where to save summaries", type=str
    )
    parser.add_argument(
        "--xsum-kwargs",
        action="store_true",
        default=False,
        help="if true use XSUM_KWARGS else CNN_KWARGS",
    )
    args = parser.parse_args()
    # main(args)
    N_process = 6
    split_data(args,N_process)
    multi_process(args,N_process)
    print("time cost",time.time()-start_time,"s")
