Create mian.py

mian.py

@@ -0,0 +1,83 @@
+!pip install -q transformers
+
+from transformers import RemBertForTokenClassification, RemBertTokenizerFast
+from transformers import XLMRobertaForTokenClassification, XLMRobertaTokenizerFast
+import torch
+
+
+
+main_path = "Misha24-10/MultiCoNER-2-recognition-model"
+
+model_1 = XLMRobertaForTokenClassification.from_pretrained(main_path,
+                                                        subfolder = "xlm_roberta_large_mountain")
+tokenizer_1 = XLMRobertaTokenizerFast.from_pretrained(main_path,
+                                                        subfolder = "xlm_roberta_large_mountain")
+
+model_2 = RemBertForTokenClassification.from_pretrained(main_path,
+                                                           subfolder = "google-rembert-ft_for_multi_ner_v3")
+tokenizer_2 = RemBertTokenizerFast.from_pretrained(main_path,
+                                                      subfolder = "google-rembert-ft_for_multi_ner_v3")
+
+model_3 = RemBertForTokenClassification.from_pretrained(main_path,
+                                                           subfolder = "google-rembert-ft_for_multi_ner_sky")
+tokenizer_3 = RemBertTokenizerFast.from_pretrained(main_path,
+                                                           subfolder = "google-rembert-ft_for_multi_ner_sky")
+
+import torch
+
+def compute_last_leyer_probs(model, tokenizer, sentence):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    number_of_tokens = tokenizer.encode_plus(sentence, return_tensors='pt',)['input_ids'].shape[-1]
+    list_of_words = sentence.split()
+
+    inputs = tokenizer(list_of_words, is_split_into_words=True, padding='max_length', max_length = min(number_of_tokens,512), truncation=True, return_tensors="pt")
+    input_ids = inputs['input_ids'].to(device)
+    attention_mask = inputs['attention_mask'].to(device)
+    label_ids = torch.Tensor(align_word_ids(inputs.word_ids()))
+    with torch.no_grad():
+        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+        logits = outputs.logits
+        return (logits[:, (label_ids == 1), :])
+
+weights = {'model_1': 1, 'model_2': 1, 'model_3': 1}
+
+def align_word_ids(word_ids, return_word_ids=False):    
+    previous_word_idx = None
+    label_ids = []
+    index_list = []
+    for idx, word_idx in enumerate(word_ids):
+
+        if word_idx is None:
+            label_ids.append(-100)
+
+        elif word_idx != previous_word_idx:
+            try:
+                label_ids.append(1)
+                index_list.append(idx)
+            except:
+                label_ids.append(-100)
+        else:
+            try:
+                label_ids.append(1 if label_all_tokens else -100)
+            except:
+                label_ids.append(-100)
+        previous_word_idx = word_idx
+
+    if return_word_ids:
+        return label_ids, index_list
+    else:
+        return label_ids
+
+def weighted_voting(sentence):
+    predictions = []
+    for idx, (model, tokenizer) in enumerate([(model_1, tokenizer_1), (model_2, tokenizer_2), (model_3, tokenizer_3)]):
+        logits = compute_last_leyer_probs(model, tokenizer, sentence)
+        predictions.append(logits * weights[f'model_{idx+1}'])
+    final_logits = sum(predictions)
+    final_predictions = torch.argmax(final_logits, dim=2)
+    labels = [model_1.config.id2label[i] for i in final_predictions.tolist()[0]]
+    return labels
+
+sent_ex = "Elon Musk 's brother sits on the boards of tesla".lower()
+
+weighted_voting(sent_ex)