Upload classifier.py

classifier.py

@@ -0,0 +1,103 @@
+from collections import defaultdict
+from transformers import BertTokenizerFast, BertForTokenClassification
+import streamlit as st
+from transformers import CharSpan
+import torch
+
+COLOURS = {
+    "ADJ": '#1DB5D7',
+    "ART": '#47A455',
+    "ELI": '#FF5252',
+    "FIN": '#9873E6',
+    "NEG": '#FF914D'
+}
+
+@st.cache(allow_output_mutation=True)
+def load_model():
+    print('Loading classification model')
+    return BertForTokenClassification.from_pretrained("classifier/")
+
+
+@st.cache(allow_output_mutation=True)
+def load_tokenizer():
+    print("Loading tokenizer for classification model")
+    return BertTokenizerFast.from_pretrained("classifier/")
+
+
+model = load_model()
+tokenizer = load_tokenizer()
+
+special_tokens = [tokenizer.unk_token, tokenizer.sep_token, tokenizer.pad_token, tokenizer.cls_token, tokenizer.mask_token]
+
+
+def predict(logits, threshold=0.3):
+  """Takes in output of model for a sentence as entry, returns one hot encoded prediction"""
+  probabilities = torch.nn.Softmax(dim=2)(logits)
+  predicted = (probabilities > threshold).long()
+  # if O is among the predicted labels, do not return other possible error types
+  for i,sent in enumerate(predicted):
+    for j,tok in enumerate(sent):
+      if tok[0] == 1:
+        predicted[i][j] = torch.tensor([1] + [0]*(model.num_labels - 1))
+  return predicted
+
+def one_hot_label_decode(tok_labels):
+  labels = []
+  for i,v in enumerate(tok_labels):
+    if v != 0:
+      labels.append(model.config.id2label[i])
+  return labels
+
+def merge_char_spans(charspans: list):
+  sorted_charspans = sorted(charspans)
+  # make groups of continuous spans
+  start = min((span.start for span in sorted_charspans))
+  end = max((span.end for span in sorted_charspans))
+  return CharSpan(start,end)
+
+def merge_error_spans(errors):
+    merged_spans = defaultdict(list)
+    for error_type in errors:
+        merged_spans[error_type] = errors[error_type][:1]
+        for charspan in errors[error_type][1:]:
+            if merged_spans[error_type][-1].end == charspan.start or merged_spans[error_type][-1].end +1 == charspan.start:
+                merged_spans[error_type][-1] = merge_char_spans([merged_spans[error_type][-1], charspan])
+    return merged_spans
+
+def predict_error_spans(sentence):
+    encoded_orig = tokenizer(sentence, return_tensors="pt")
+    tokens = tokenizer.convert_ids_to_tokens(encoded_orig.input_ids[0])
+    output = model(**encoded_orig)
+    predicted_labels = predict(output.logits, threshold=0.2)
+    errors = defaultdict(list)
+
+    for i,(t,p) in enumerate(list(zip(tokens, predicted_labels[0]))):
+        labels = one_hot_label_decode(p)
+        if t not in special_tokens and labels != ['O'] : 
+            for error_type in labels:
+                position, e_type = error_type.split("-")
+                if position == "B":
+                    errors[e_type].append(encoded_orig.token_to_chars(i))
+                else:
+                    errors[e_type][-1] = merge_char_spans([errors[e_type][-1], encoded_orig.token_to_chars(i)])
+    return sentence, merge_error_spans(errors)
+
+def annotate(sentence: str):
+    sent_and_error_spans = predict_error_spans(sentence)
+    error_spans_ordered = []
+    for error_type, charspans in sent_and_error_spans[1].items():
+        for charspan in charspans:
+            error_spans_ordered.append((charspan,error_type))
+        error_spans_ordered = sorted(error_spans_ordered, key= lambda x:x[0])
+    annotated_sentence = []
+    for i,(error_span,label) in enumerate(error_spans_ordered):
+        if i > 0:
+            annotated_sentence.append(sentence[error_spans_ordered[i-1][0].end:error_span.start])
+        else:
+            annotated_sentence.append(sentence[:error_span.start])
+        annotated_sentence.append((sentence[error_span.start:error_span.end], label, COLOURS[label]))
+    if len(error_spans_ordered) > 0:
+        annotated_sentence.append(sentence[error_span.end:])
+    else:
+        annotated_sentence.append(sentence)
+    return annotated_sentence