| from transformers import Pipeline | |
| import tensorflow as tf | |
| 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 json | |
| import numpy as np | |
| class NewsClassifierPipeline(Pipeline): | |
| def __init__(self): | |
| super().__init__() | |
| self.model = load_model('./news_classifier.h5') | |
| with open('tokenizer.json') as f: | |
| tokenizer_data = json.load(f) | |
| self.tokenizer = tokenizer_from_json(tokenizer_data) | |
| def preprocess(self, text): | |
| sequence = self.tokenizer.texts_to_sequences([text]) | |
| padded = pad_sequences(sequence) | |
| return padded | |
| def _forward(self, texts): | |
| processed = self.preprocess(texts) | |
| predictions = self.model.predict(processed) | |
| scores = tf.nn.softmax(predictions, axis=1) | |
| predicted_class = np.argmax(predictions) | |
| score = float(np.max(scores)) | |
| label = 'foxnews' if predicted_class == 0 else 'nbc' | |
| return [{'label': label, 'score': score}] | |
| def postprocess(self, model_outputs): | |
| return model_outputs |