diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e0d18a6 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ +from .classifier import ClassifierNode + +NODE_CLASS_MAPPINGS = { + "ClassifierNode": ClassifierNode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ClassifierNode": "Classifier Node" +} \ No newline at end of file diff --git a/classifier.py b/classifier.py new file mode 100644 index 0000000..86928a9 --- /dev/null +++ b/classifier.py @@ -0,0 +1,48 @@ +import json +from transformers import pipeline + +class ClassifierNode: + def __init__(self): + self.model = pipeline("zero-shot-classification") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text_input": ("STRING", {"default": ""}), + "candidate_labels_json": ("STRING", {"default": "{\"tags\": []}"}), + "num_labels": ("INT", {"default": 10}), + "include_scores": ("BOOLEAN", {"default": True}), + "output_format": (["json", "text"], {"default": "json"}) + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("classification_result",) + FUNCTION = "classify" + CATEGORY = "Classifier" + + def classify(self, text_input, candidate_labels_json, num_labels=10, include_scores=True, output_format="json"): + try: + # Parse candidate labels from JSON + candidate_labels_data = json.loads(candidate_labels_json) + labels = [item['tag'] for item in candidate_labels_data['tags']] + + # Perform zero-shot classification + result = self.model(text_input, labels) + + # Prepare the output + output = [] + for label, score in zip(result['labels'], result['scores']): + if len(output) < num_labels: + if include_scores: + output.append({"tag": label, "score": score}) + else: + output.append({"tag": label}) + + if output_format == 'json': + return (json.dumps({"tags": output}),) + else: + return (str({"tags": output}),) + except Exception as e: + return (json.dumps({"error": str(e)}),) \ No newline at end of file diff --git a/demo.png b/demo.png new file mode 100644 index 0000000..6829bdc Binary files /dev/null and b/demo.png differ