29 lines
916 B
Python
29 lines
916 B
Python
from transformers import pipeline as hf_pipeline
|
|
|
|
text_classification_model_list = [
|
|
"distilbert-base-uncased-finetuned-sst-2-english",
|
|
"cardiffnlp/twitter-roberta-base-sentiment-latest",
|
|
]
|
|
|
|
|
|
class TextClassificationPipeline:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"default": "", "multiline": True}),
|
|
"model_name": (text_classification_model_list, {"default": text_classification_model_list[0]}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "FLOAT",)
|
|
RETURN_NAMES = ("label", "score",)
|
|
FUNCTION = "run_classification"
|
|
CATEGORY = "Transformers/NLP/TextClassification"
|
|
|
|
def run_classification(self, text, model_name):
|
|
pipe = hf_pipeline("text-classification", model=model_name)
|
|
result = pipe(text)
|
|
top = result[0]
|
|
return (top["label"], top["score"],)
|