from datasets import load_dataset
from SBERT_Model import SBERT_Model

test_dataset = load_dataset('paws', 'labeled_final', split='test')
test_dataset.set_format(type='pandas')
test_dataset = test_dataset[:]

zero_model = SBERT_Model("Zero Model", '../models/nli-bert-large/', test_dataset)
ft_model = SBERT_Model("Fine-tuned Model", '../models/stsb-bert-large/', test_dataset)
paws_ft_model = SBERT_Model("Fine-tuned on PAWS Model", "models/paws-stsb-bert-large", test_dataset)

zero_model.print_statistics()
ft_model.print_statistics()
paws_ft_model.print_statistics()

zero_model.visualize_preds()
ft_model.visualize_preds()
paws_ft_model.visualize_preds()
