from transformers import Pipeline
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.text import tokenizer_from_json
from tensorflow.keras.preprocessing.sequence import pad_sequences
import numpy as np
import tensorflow as tf
import json

class NewsClassifierPipeline(Pipeline):
    def __init__(self, model_path="news_classifier.h5", tokenizer_path="tokenizer.json"):
        super().__init__()
        self.model = load_model(model_path)
        with open(tokenizer_path, "r") as f:
            tokenizer_data = json.load(f)
            self.tokenizer = tokenizer_from_json(tokenizer_data)

    def preprocess(self, inputs):
        sequences = self.tokenizer.texts_to_sequences([inputs])
        return pad_sequences(sequences, maxlen=128)

    def _forward(self, inputs):
        preprocessed = self.preprocess(inputs)
        predictions = self.model.predict(preprocessed)
        scores = tf.nn.softmax(predictions, axis=1).numpy()
        label = np.argmax(scores)
        return [{"label": "foxnews" if label == 0 else "nbc", "score": float(scores[0, label])}]

    def postprocess(self, model_outputs):
        return model_outputs