code/do_predict.py

from args import args, config
from items_dataset import items_dataset
from torch.utils.data import DataLoader
from models import Model_Crf, Model_Softmax
from transformers import AutoTokenizer
from tqdm import tqdm
import prediction
import torch
import math

directory = args.SAVE_MODEL_PATH
model_name = "roberta_CRF.pt"
device = torch.device("cuda", 0) if torch.cuda.is_available() else torch.device("cpu")
model_crf = Model_Crf(config).to(device)
model_crf.load_state_dict(
    state_dict=torch.load(directory + model_name, map_location=device)
)

model_name = "roberta_softmax.pt"
device = torch.device("cuda", 0) if torch.cuda.is_available() else torch.device("cpu")
model_roberta = Model_Softmax(config).to(device)
model_roberta.load_state_dict(
    state_dict=torch.load(directory + model_name, map_location=device)
)


def prepare_span_data(dataset):
    for sample in dataset:
        spans = items_dataset.cal_agreement_span(
            None,
            agreement_table=sample["predict_sentence_table"],
            min_agree=1,
            max_agree=2,
        )
        sample["span_labels"] = spans
        sample["original_text"] = sample["text_a"]
        del sample["text_a"]


def rank_spans(test_loader, device, model, reverse=True):
    """Calculate each span probability by e**(word average log likelihood)"""
    model.eval()
    result = []

    for i, test_batch in enumerate(tqdm(test_loader)):
        batch_text = test_batch["batch_text"]
        input_ids = test_batch["input_ids"].to(device)
        token_type_ids = test_batch["token_type_ids"].to(device)
        attention_mask = test_batch["attention_mask"].to(device)
        labels = test_batch["labels"]
        crf_mask = test_batch["crf_mask"].to(device)
        sample_mapping = test_batch["overflow_to_sample_mapping"]
        output = model(
            input_ids=input_ids,
            token_type_ids=token_type_ids,
            attention_mask=attention_mask,
            labels=None,
            crf_mask=crf_mask,
        )
        output = torch.nn.functional.softmax(output[0], dim=-1)

        # make result of every sample
        sample_id = 0
        sample_result = {
            "original_text": test_batch["batch_text"][sample_id],
            "span_ranked": [],
        }
        for batch_id in range(len(sample_mapping)):
            change_sample = False

            # make sure status
            if sample_id != sample_mapping[batch_id]:
                change_sample = True
            if change_sample:
                sample_id = sample_mapping[batch_id]
                result.append(sample_result)
                sample_result = {
                    "original_text": test_batch["batch_text"][sample_id],
                    "span_ranked": [],
                }

            encoded_spans = items_dataset.cal_agreement_span(
                None, agreement_table=labels[batch_id], min_agree=1, max_agree=2
            )
            # print(encoded_spans)
            for encoded_span in encoded_spans:
                # calculate span loss
                span_lenght = encoded_span[1] - encoded_span[0]
                # print(span_lenght)
                span_prob_table = torch.log(
                    output[batch_id][encoded_span[0] : encoded_span[1]]
                )
                if (
                    not change_sample and encoded_span[0] == 0 and batch_id != 0
                ):  # span cross two tensors
                    span_loss += span_prob_table[0][1]  # Begin
                else:
                    span_loss = span_prob_table[0][1]  # Begin
                for token_id in range(1, span_prob_table.shape[0]):
                    span_loss += span_prob_table[token_id][2]  # Inside
                span_loss /= span_lenght

                # span decode
                decode_start = test_batch[batch_id].token_to_chars(encoded_span[0] + 1)[
                    0
                ]
                decode_end = test_batch[batch_id].token_to_chars(encoded_span[1])[0] + 1
                # print((decode_start, decode_end))
                span_text = test_batch["batch_text"][sample_mapping[batch_id]][
                    decode_start:decode_end
                ]
                if (
                    not change_sample and encoded_span[0] == 0 and batch_id != 0
                ):  # span cross two tensors
                    presample = sample_result["span_ranked"].pop(-1)
                    sample_result["span_ranked"].append(
                        [presample[0] + span_text, math.e ** float(span_loss)]
                    )
                else:
                    sample_result["span_ranked"].append(
                        [span_text, math.e ** float(span_loss)]
                    )
        result.append(sample_result)

    # sorted spans by probability
    # for sample in result:
    #     sample["span_ranked"] = sorted(
    #         sample["span_ranked"], key=lambda x: x[1], reverse=reverse
    #     )
    return result


def predict_single(text):
    input_dict = [{"span_labels": []}]
    input_dict[0]["original_text"] = text
    tokenizer = AutoTokenizer.from_pretrained(
        args.pre_model_name, add_prefix_space=True
    )
    prediction_dataset = items_dataset(tokenizer, input_dict, args.label_dict)
    prediction_loader = DataLoader(
        prediction_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=prediction_dataset.collate_fn,
    )
    predict_data = prediction.test_predict(prediction_loader, device, model_crf)
    prediction.add_sentence_table(predict_data)

    prepare_span_data(predict_data)
    tokenizer = AutoTokenizer.from_pretrained(
        args.pre_model_name, add_prefix_space=True
    )
    prediction_dataset = items_dataset(tokenizer, predict_data, args.label_dict)
    prediction_loader = DataLoader(
        prediction_dataset,
        batch_size=args.batch_size,
        shuffle=False,
        collate_fn=prediction_dataset.collate_fn,
    )
    span_ranked = rank_spans(prediction_loader, device, model_roberta)
    # for sample in span_ranked:
    #     print(sample["original_text"])
    #     print(sample["span_ranked"])

    result = []
    sample = span_ranked[0]
    orig = sample["original_text"]
    cur = 0
    for s, score in sample["span_ranked"]:
        # print()
        # print('ORIG', repr(orig))
        # print('CCUR', repr(orig[cur:]))
        # print('SSSS', repr(s))
        # print()
        end = orig.index(s, cur)
        if cur != end:
            result.append([orig[cur:end], 0])
        result.append([s, score])
        cur = end + len(s)
    if cur < len(orig):
        result.append([orig[cur:], 0])
    return result


if __name__ == "__main__":
    s = """貓咪犯錯後，以下5種懲罰方法很有效，飼主可以試試樂享網 2021-03-06 繼續閱讀 繼續閱讀 繼續閱讀 繼續閱讀 繼續閱讀 貓咪雖然高冷，但也是會犯錯的，那貓咪犯錯後，怎麼懲罰它才最有效呢？今天就來說一些懲罰貓咪最有效的5個方法！1、把痛感形成條件反射 這裡說的是「痛感」，而不是「暴打」。在貓咪犯錯後，寵主不需要打它，可以彈鼻頭或者是輕拍它的頭頂，給它造成痛感，這樣讓貓咪有一些畏懼心理，知道你在懲罰它。這樣時間長了，貓咪就會形成條件反射，以後就會少犯錯了。  2、大聲呵斥比起打貓，大聲呵斥貓咪會更加有效。因為貓咪對聲音很敏感，它能從主人的語氣中判斷主人的情緒，當大聲呵斥它的時候，它往往會楞一下，這時你繼續大聲呵斥它，那它就會明白你不允許它做這件事，這樣犯錯地方幾率就會減少了。  3、限制自由限制自由說白了，就是把貓咪關進籠子裡。因為貓咪都是很愛外出玩耍，當它犯錯咯，主人可以把它關進籠子裡，不搭理它，讓它自己反思。但要注意，這個方法不能經常用，而且不能把貓咪關進籠子太久。  4、利用水都知道貓咪很怕水的，所以當貓咪犯錯後，寵主也可以利用水來懲罰貓咪，這也是很效果的方法。寵主可以給貓咪臉上或是頭頂噴一些水，從而讓貓知道這個行為是錯誤的，以後就不會再犯了。  5、冷淡雖然貓咪不是很粘主人，但它還是很愛主人的，所以在貓咪犯錯後，寵主也可以採取冷淡的方法來懲罰貓。對貓咪採取不理睬、不靠近、不擁抱等策略，這樣貓咪就會知道自己錯了。當然懲罰的時間不要太長，不然貓咪就會以為你不愛它了。"""
    print(predict_single(s))