add classifier

1 files changed, 10 insertions, 0 deletions

diff --git a/classification/main.py b/classification/main.py
new file mode 100755
index 000000000..04f2d8c49
--- /dev/null
+++ b/classification/main.py
@@ -0,0 +1,10 @@
+from transformers import pipeline
+
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+with open("test", "r") as file:
+    sequence_to_classify = file.read()
+candidate_labels = ['semantic bug', 'no semantic bug']
+result = classifier(sequence_to_classify, candidate_labels, multi_label=False)
+
+print(result['labels'])
+print(result['scores'])