Restructured files

_init_.py

@@ -1,4 +0,0 @@
-from .modeling import TFNewsClassifier
-from transformers import AutoConfig
-
-AutoConfig.register("news_classifier", TFNewsClassifier)

config.json

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:87fdbfd9402e0152bf788f338996d8fce42dbafc705bb7c8fc7fb8f84d6817e8
-size 307
+oid sha256:81c482c90c48b46c5f596d6668d2156f438723e160ffe30628880b28fdfc38ab
+size 241

modeling.py

@@ -1,21 +0,0 @@
-import tensorflow as tf
-from transformers import PreTrainedModel, TFPreTrainedModel
-from transformers.modeling_tf_utils import TFSequenceClassificationLoss
-
-class TFNewsClassifier(TFPreTrainedModel, TFSequenceClassificationLoss):
-    def __init__(self, config):
-        super().__init__(config)
-        self.num_labels = config.num_labels
-        
-        # Define your model architecture
-        self.lstm = tf.keras.layers.LSTM(128, return_sequences=True)
-        self.lstm2 = tf.keras.layers.LSTM(64)
-        self.dropout = tf.keras.layers.Dropout(0.5)
-        self.classifier = tf.keras.layers.Dense(self.num_labels, activation='softmax')
-        
-    def call(self, inputs, training=False):
-        x = self.lstm(inputs)
-        x = self.dropout(x, training=training)
-        x = self.lstm2(x)
-        x = self.dropout(x, training=training)
-        return self.classifier(x)

pipeline.py

@@ -0,0 +1,35 @@
+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