File size: 1,094 Bytes
de435e1 7532efb de435e1 9a2d08e de435e1 9a2d08e de435e1 072da58 de435e1 7532efb 072da58 f4e5a43 de435e1 f4e5a43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | 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 json
class NewsClassifierPipeline(Pipeline):
def __init__(self):
super().__init__()
self.model = load_model('./news_classifier.h5')
with open('./tokenizer.json', '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):
processed = self.preprocess(inputs)
predictions = self.model.predict(processed)
label = "foxnews" if predictions[0][0] > 0.5 else "nbc"
score = predictions[0][0] if label == "foxnews" else 1 - predictions[0][0]
return [{"label": label, "score": float(score)}]
def postprocess(self, outputs):
return outputs
|