infer.py

import sys, pathlib

sys.path.append(pathlib.Path(__file__).parent.as_posix())

import torch
import utils.constant as constant
from transformers import (
    pipeline,
    AutoTokenizer,
    AutoModelForSequenceClassification,
)

""" Load tokenizer and model. Setup pipeline. """
tokenizer = AutoTokenizer.from_pretrained(constant.FINETUNED_MODEL)
model = AutoModelForSequenceClassification.from_pretrained(constant.FINETUNED_MODEL)
classifier = pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
)

while True:
    text = input("Enter comment to classify (or 'exit' to quit): ")
    if text.lower() == 'exit':
        break
    
    outputs = classifier(text, return_all_scores=True)
    print(outputs)