CoolFace
Modelpublic

NightPrince/Toxic_Classification

sourceHugging Facemitupdated 1y agoView on Hugging Face
0likes22downloads
pipeline.py41 linesDownload Raw Back to root
1import numpy as np2import tensorflow as tf3from tensorflow.keras.preprocessing.sequence import pad_sequences4from tensorflow.keras.preprocessing.text import tokenizer_from_json5import json6import os7 8# Hugging Face expects a class named Pipeline with __call__(self, inputs)9class Pipeline:10    def __init__(self):11        # Load tokenizer12        with open("tokenizer.json", "r", encoding="utf-8") as f:13            tokenizer_json = f.read()14            self.tokenizer = tokenizer_from_json(tokenizer_json)15        self.max_len = 15016 17        # Load model (SavedModel format)18        self.model = tf.keras.models.load_model(".")19 20        # Load label map if available21        self.label_map = None22        if os.path.exists("label_map.json"):23            with open("label_map.json", "r", encoding="utf-8") as f:24                self.label_map = json.load(f)25 26    def __call__(self, inputs):27        # Accepts a dict with keys 'text' and 'image_desc'28        text = inputs.get("text", "")29        image_desc = inputs.get("image_desc", "")30        input_text = text + " " + image_desc31        seq = self.tokenizer.texts_to_sequences([input_text])32        padded = pad_sequences(seq, maxlen=self.max_len, padding='post', truncating='post')33        pred_probs = self.model.predict(padded)34        pred_label = int(np.argmax(pred_probs, axis=1)[0])35        score = float(np.max(pred_probs))36        if self.label_map:37            label = self.label_map.get(str(pred_label), pred_label)38        else:39            label = pred_label40        return {"label": label, "score": score}41