import unittest from models.full_finetune_model import train_full_finetune_model class TestModels(unittest.TestCase): def test_train_full_finetune_model(self): train_data = preprocess_data(load_dataset('data/train.csv')) val_data = preprocess_data(load_dataset('data/val.csv')) model = train_full_finetune_model(train_data, val_data) self.assertIsNotNone(model) if __name__ == '__main__': unittest.main()