sentiment / scripts /test.py
danielle2003's picture
adding scripts'
2b3de4f
raw
history blame contribute delete
544 Bytes
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()