predict.py

from utils import *
from dataloader import *
from rnn_crf import *

def load_model(args):

    cti = load_tkn_to_idx(args[1]) # char_to_idx
    wti = load_tkn_to_idx(args[2]) # word_to_idx
    itt = load_idx_to_tkn(args[3]) # idx_to_tag

    model = rnn_crf(cti, wti, len(itt))
    print(model)

    load_checkpoint(args[0], model)

    return model, cti, wti, itt

def run_model(model, data, itt):

    with torch.no_grad():
        model.eval()

        for batch in data.batchify(BATCH_SIZE):

            xc, xw, lens = batch.xc, batch.xw, batch.lens
            xc, xw = data.to_tensor(bc = xc, bw = xw, lens = lens)
            y1 = model.decode(xc, xw, lens)
            batch.y1 = [[itt[i] for i in y] for y in y1]

            for x0, y0, y1 in zip(batch.x0, batch.y0, batch.y1):
                if not HRE:
                    x0, y0, y1 = [x0], [y0], [y1]
                for x0, y0, y1 in zip(x0, y0, y1):
                    yield x0, y0, y1

def predict(model, cti, wti, itt, filename):

    data = dataloader(hre = HRE)
    with open(filename) as fo:
        text = fo.read().strip().split("\n" * (HRE + 1))

    for block in text:
        data.append_row()

        for line in block.split("\n"):

            if re.match("\S+/[^ /]+( \S+/[^ /]+)*$", line): # word/tag
                x0, y0 = zip(*[re.split("/(?=[^/]+$)", w) for w in line.split(" ")])
                x1 = list(map(normalize, x0))
            else:
                x0, y0 = line, []
                if re.match("[^\t]+\t[^\t]+$", x0): # sentence \t label
                    x0, *y0 = x0.split("\t")
                x0 = tokenize(x0)
                x1 = list(map(normalize, x0))

            xc = [[cti[c] if c in cti else UNK_IDX for c in w] for w in x1]
            xw = [wti[w] if w in wti else UNK_IDX for w in x1]

            data.append_item(x0 = x0, x1 = x1, xc = xc, xw = xw, y0 = y0)

    return run_model(model, data, itt)

if __name__ == "__main__":

    if len(sys.argv) != 6:
        sys.exit("Usage: %s model char_to_idx word_to_idx tag_to_idx test_data" % sys.argv[0])

    result = predict(*load_model(sys.argv[1:5]), sys.argv[5])
    func = tag_to_txt if TASK else lambda *x: x

    for x0, y0, y1 in result:
        if y0:
            print(func(x0, y0))
        print(func(x0, y1))