| from transformers import Pipeline |
| import tensorflow as tf |
| from tensorflow.keras.preprocessing.sequence import pad_sequences |
| import json |
| import os |
|
|
| def load_tokenizer(tokenizer_path): |
| with open(tokenizer_path, 'r') as f: |
| return json.load(f) |
|
|
| class NewsClassificationPipeline(Pipeline): |
| def __init__(self, model=None, tokenizer=None, **kwargs): |
| super().__init__(**kwargs) |
| model_path = os.path.join(os.path.dirname(__file__), 'news_classifier.h5') |
| self.model = tf.keras.models.load_model(model_path) |
|
|
| tokenizer_path = os.path.join(os.path.dirname(__file__), 'tokenizer.json') |
| self.tokenizer_config = load_tokenizer(tokenizer_path) |
| |
| def __call__(self, texts, **kwargs): |
| if isinstance(texts, str): |
| texts = [texts] |
|
|
| sequences = self.tokenizer.texts_to_sequences(texts) |
| padded = pad_sequences(sequences, maxlen=128) |
| |
| predictions = self.model.predict(padded) |
| |
| results = [] |
| for pred in predictions: |
| label = "foxnews" if pred[0] > 0.5 else "nbc" |
| score = float(pred[0] if label == "foxnews" else 1 - pred[0]) |
| results.append({"label": label, "score": score}) |
| |
| return results[0] if isinstance(texts, str) else results |
