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()