!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)