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))