Create model_tools.py

model_tools.py

@@ -0,0 +1,99 @@
+import re
+import torch
+from torch import cuda
+from nltk.tokenize.casual import TweetTokenizer
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+pretrained_name = "ml6team/xlm-roberta-base-nl-emoji-ner"
+tokenizer = AutoTokenizer.from_pretrained(pretrained_name)
+model = AutoModelForTokenClassification.from_pretrained(pretrained_name)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+emoji_list = ['😨', '😥', '😍', '😠', '🤯', '😄', '🍾', '🚗', '☕', '💰']
+emoji_names = ['B-AFRAID', 'B-SAD', 'B-LOVE', 'B-ANGRY', 'B-SHOCKED', 'B-LAUGH', 'B-CHAMP', 'B-CAR', 'B-COFFEE', 'B-MONEY']
+
+from collections import defaultdict
+
+def default_value():
+    return 0
+
+emoji2tag = {emo:tag for emo, tag in zip(emoji_list, emoji_names)}
+emoji2tag['O'] = 'O'
+tag2emoji = {val:key for key,val in emoji2tag.items()}
+
+tag2idx = defaultdict(default_value)
+
+# mapping of emojis and categories
+for i, e in enumerate(emoji_names):
+    tag2idx[e] = i+1
+
+tag2idx['O'] = 0
+
+
+idx2tag = {val:key for key,val in tag2idx.items()}
+
+
+def tag_text_sample(text):
+
+    t = TweetTokenizer()
+
+    words = t.tokenize(text)
+    # Get tokens with special characters
+    tokens = tokenizer(words, is_split_into_words=True).tokens()
+
+    # Encode the sequence into IDs
+    input_ids = tokenizer(
+        words,
+        add_special_tokens=True, 
+        is_split_into_words=True,
+        return_tensors="pt").input_ids.to(device)
+
+    # Get predictions as distribution over 7 possible classes
+    model.eval()
+    outputs = model(input_ids)[0]
+
+    # Take argmax to get most likely class per token
+    predictions = torch.argmax(outputs, dim=2)
+
+    # Convert to DataFrame
+    preds = [idx2tag[p] for p in predictions[0].cpu().numpy()]
+
+    # Get word id's
+    word_ids = tokenizer(words, is_split_into_words=True).word_ids()
+
+    # Keep non-special tokens and labels
+    preds = preds[1:-1]
+    tokens = tokens[1:-1]
+    word_ids = word_ids[1:-1]
+
+    # Full sentence reconstruction
+    previous_word_idx = None
+    current_emoji = ''
+
+    word_emojis = {}
+
+    for i, (word_idx, pred) in enumerate(zip(word_ids, preds)):
+
+        if previous_word_idx is None or word_idx != previous_word_idx:
+            # Add emoji from previous word to sequence
+            word_emojis[word_idx] = current_emoji
+
+            # Check new emoji
+            current_emoji = tag2emoji[pred] if not pred == "O" else ""
+
+        previous_word_idx = word_idx
+
+        # If final word: add emoji
+        if i == len(word_ids) - 1:
+            word_emojis[word_idx] = current_emoji
+
+    # Reconstruct
+    full_text = []
+    for i, word in enumerate(words):
+
+        full_text.append(word)
+
+        if word_emojis[i]:
+            full_text.append(word_emojis[i])
+
+    return tokenizer.clean_up_tokenization(" ".join(full_text))