Upload get_predictions.py

get_predictions.py

@@ -0,0 +1,60 @@
+import tokenize_uk
+import torch
+
+def get_word_predictions(model, tokenizer, texts, is_split_to_words=False, device='cpu'):
+    words_res = []
+    y_res = []
+
+    if not is_split_to_words:
+        texts = [tokenize_uk.tokenize_words(text) for text in texts]
+
+    for text in texts:
+        size = len(text)
+        idx_list = [idx + 1 for idx, val in enumerate(text) if val in ['.', '?', '!']]
+        if len(idx_list):
+            sents = [text[i: j] for i, j in zip([0] + idx_list, idx_list + ([size] if idx_list[-1] != size else []))]
+        else:
+            sents = [text]
+
+        y_res_x = []
+        words_res_x = []
+        for sent_tokens in sents:
+            tokenized_inputs = [101]
+            word_ids = [None]
+            for word_id, word in enumerate(sent_tokens):
+                word_tokens = tokenizer.encode(word)[1:-1]
+                tokenized_inputs += word_tokens
+                word_ids += [word_id]*len(word_tokens)
+            tokenized_inputs = tokenized_inputs[:(tokenizer.model_max_length-1)]
+            word_ids = word_ids[:(tokenizer.model_max_length-1)]
+            tokenized_inputs += [102]
+            word_ids += [None]
+
+            torch_tokenized_inputs = torch.tensor(tokenized_inputs).unsqueeze(0)
+            torch_attention_mask = torch.ones(torch_tokenized_inputs.shape)
+            predictions = model.forward(input_ids=torch_tokenized_inputs.to(device), attention_mask=torch_attention_mask.to(device))
+            predictions = torch.argmax(predictions.logits.squeeze(), axis=1).numpy()
+            predictions = [model.config.id2label[i] for i in predictions]
+
+            previous_word_idx = None
+            sent_words = []
+            predictions_words = []
+            word_tokens = []
+            first_pred = None
+            for i, word_idx in enumerate(word_ids):
+                if word_idx != previous_word_idx:
+                    sent_words.append(tokenizer.decode(word_tokens))
+                    word_tokens = [tokenized_inputs[i]]
+                    predictions_words.append(first_pred)
+                    first_pred = predictions[i]
+                else:
+                    word_tokens.append(tokenized_inputs[i])      
+                previous_word_idx = word_idx
+
+            words_res_x.extend(sent_words[1:])
+            y_res_x.extend(predictions_words[1:])
+
+        words_res.append(words_res_x)
+        y_res.append(y_res_x)
+
+    return words_res, y_res