File size: 3,129 Bytes
3b434b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np
import torch
model = AutoModelForSequenceClassification.from_pretrained("DavinciTech/BERT_Categorizer")
tokenizer = AutoTokenizer.from_pretrained("DavinciTech/BERT_Categorizer")
model.to("cuda")
input_texts = [
("Title: Scanner not working\n"
"Description: Good morning Team My scanner is not connecting to the image saving folder I assume it has "
"something to do with the merge from last week I need to scan to be able to do payroll which needs to all be done "
"before and scanning is the first step"),
("Title: New mouse please\n"
"Description: My mouse is acting up a little bit could I get a new one please?"),
("Title: Internet outage\n"
"Description: The whole internet is down for everyone in the office"),
]
id2label = {k: l for k, l in enumerate(model.config.LABEL_DICTIONARY.keys())}
#label2id = {l: k for k, l in enumerate(model.config.LABEL_DICTIONARY.keys())}
# Encode the text
encoded = tokenizer(input_texts, truncation=True, padding="max_length", max_length=512, return_tensors="pt").to("cuda")
# Call the model to predict under the format of logits of 27 classes
logits = model(**encoded).logits.cpu().detach().numpy()
IMPACT_LABELS = ["I1", "I2", "I3", "I4"]
IMPACT_INDICES = range(0, 4)
URGENCY_LABELS = ["U1", "U2", "U3", "U4"]
URGENCY_INDICES = range(4, 8)
TYPE_LABELS = ["T1", "T2", "T3", "T4", "T5"]
TYPE_INDICES = range(8, 13)
ALL_LABELS = IMPACT_LABELS + URGENCY_LABELS + TYPE_LABELS
def get_preds_from_logits(logits):
ret = np.zeros(logits.shape)
# The first 5 columns (IMPACT_INDICES) are for Impact. They should be handled with a multiclass approach
# i.e. we fill 1 to the class with highest probability, and 0 into the other columns
best_class = np.argmax(logits[:, IMPACT_INDICES], axis=-1)
ret[list(range(len(ret))), best_class] = 1
ret[:, URGENCY_INDICES] = 0 # Initialize all priority indices to 0
ret[:, TYPE_INDICES] = 0 # Initialize all type indices to 0
# Find the index with the maximum value in the PRIORITY_INDICES and set it to 1
max_priority_index = np.argmax(logits[:, URGENCY_INDICES], axis=-1)
ret[list(range(len(ret))), max_priority_index + URGENCY_INDICES[0]] = 1
# Find the index with the maximum value in the TYPE_INDICES and set it to 1
max_type_index = np.argmax(logits[:, TYPE_INDICES], axis=-1)
ret[list(range(len(ret))), max_type_index + TYPE_INDICES[0]] = 1
return ret
# Decode the result
preds = get_preds_from_logits(logits)
decoded_preds = [[id2label[i] for i, l in enumerate(row) if l == 1] for row in preds]
print("\n")
for text, pred in zip(input_texts, decoded_preds):
print(text)
print("Impact:", [model.config.LABEL_DICTIONARY[l] for l in pred if l.startswith("I")])
print("Urgency:", [model.config.LABEL_DICTIONARY[l] for l in pred if l.startswith("U")])
print("Type:", [model.config.LABEL_DICTIONARY[l] for l in pred if l.startswith("T")])
print("")
|