sts_pair_explainer.py

from nltk.tokenize import word_tokenize
import pandas as pd
import numpy as np
import shap
import os
import shutil
import subprocess

from sts_models.psbert_model import PSBertModel
from sts_models.sbert_model import SBertModel
from sts_models.bertscore_model import BertScoreModel

class STSWrapper():
    def __init__(self, sts_model, tokenizer='nltk'):
        self.sts_model = sts_model
        if tokenizer == 'nltk': self.tokenizer = word_tokenize
        elif tokenizer == 'split': self.tokenizer = self._split_tokenizer
            
    def _split_tokenizer(self, sent):
        return sent.split()

    def __call__(self, sent_pair_list):
        batch = []
        for pair in sent_pair_list:
            s1,s2 = pair[0].split('[SEP]')
            batch.append( (s1,s2) )

        scores = self.sts_model(batch)
        return scores

    def _tokenize_sent(self, sentence):
        if isinstance(sentence,str):
            #token_ids = self.sts_model.tokenizer.encode(sentence)
            #tokens = self.sts_model.tokenizer.convert_ids_to_tokens(token_ids)[1:-1]
            tokens = self.tokenizer(sentence)
        elif isinstance(sentence, list):
            tokens = sentence

        return tokens

    def build_feature(self, sent1, sent2):
        tokens1 = self._tokenize_sent(sent1)
        tokens2 = self._tokenize_sent(sent2)
        self.s1len = len(tokens1)
        self.s2len = len(tokens2)

        tdict = {}
        for i in range(len(tokens1)):
            tdict['s1_{}'.format(i)] = tokens1[i]
        for i in range(len(tokens2)):
            tdict['s2_{}'.format(i)] = tokens2[i]

        return pd.DataFrame(tdict, index=[0])

    def mask_model(self, mask, x):
        tokens = []
        for mm, tt in zip(mask, x):
            if mm: tokens.append(tt)
            else: tokens.append('[MASK]')
        s1 = ' '.join(tokens[:self.s1len])
        s2 = ' '.join(tokens[self.s1len:])
        sentence_pair = pd.DataFrame([s1+'[SEP]'+s2])
        return sentence_pair



class ExplainableSTS():
    def __init__(self, wanted_sts_model):
        if wanted_sts_model == 'sbert':
            sts_model = SBertModel()
        elif wanted_sts_model == 'pair-bert':
            sts_model = PSBertModel()
        elif wanted_sts_model == 'bert-score':
            sts_model = BertScoreModel()

        self.wrapper = STSWrapper(sts_model)

    def __call__(self, sent1, sent2):
        s1 = ' '.join(self.wrapper.tokenizer(sent1))
        s2 = ' '.join(self.wrapper.tokenizer(sent2))
        return self.wrapper.sts_model([(s1, s2)])[0]

    def explain(self, sent1, sent2, plot=False):
        explainer = shap.Explainer(self.wrapper, self.wrapper.mask_model)
        value = explainer(self.wrapper.build_feature(sent1, sent2))
        if plot: shap.waterfall_plot(value[0])
        all_tokens = [] 
        all_tokens += ['s1_'+t for t in self.wrapper.tokenizer(sent1)] 
        all_tokens += ['s2_'+t for t in self.wrapper.tokenizer(sent2)] 

        return [(token,sv) for token, sv in zip(all_tokens,value[0].values)]