from src.SBERT_Model import SBERT_Model
from STSB_Dataset import STSB_Dataset
# import probing # when needed

dataset = STSB_Dataset('../data/stsbenchmark/sts-test.csv').as_dataframe

zero_model = SBERT_Model("nli-bert-large", 'models/nli-bert-large/', dataset)
ft_model = SBERT_Model("stsb-bert-large", '../models/stsb-bert-large/', dataset)
paws_model = SBERT_Model("paws-stsb-bert-large", '../models/stsb-bert-large/', dataset)

zero_model.print_statistics()
ft_model.print_statistics()
paws_model.print_statistics()


#Make a list with sentence pairs for the heatmap
#sentence_pairs = []
#for i in range(2):
#    sentence_pairs.append((ft_model.sentences1[i],ft_model.sentences2[i]))
#generate the heatmap with default mask == "[MASK]"
#ft_model.get_html_word_importance(sentence_pairs)
