Add custom model architecture and configuration

_init_.py

@@ -0,0 +1,4 @@
+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:81c482c90c48b46c5f596d6668d2156f438723e160ffe30628880b28fdfc38ab
-size 241
+oid sha256:87fdbfd9402e0152bf788f338996d8fce42dbafc705bb7c8fc7fb8f84d6817e8
+size 307

modeling.py

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

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