Spaces:
Sleeping
Sleeping
File size: 544 Bytes
2b3de4f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
import unittest
from transformers import pipeline
classifier = pipeline("text-classification", model="./models")
class TestModel(unittest.TestCase):
def test_positive_sentiment(self):
result = classifier("I love this product!")[0]
self.assertIn(result["label"], ["LABEL_0", "LABEL_1", "LABEL_2"])
def test_negative_sentiment(self):
result = classifier("This is terrible, I hate it.")[0]
self.assertIn(result["label"], ["LABEL_0", "LABEL_1", "LABEL_2"])
if __name__ == "__main__":
unittest.main()
|