Updated pipeline

pipeline.py

@@ -1,33 +1,31 @@
-from transformers import Pipeline
+from transformers import PreTrainedModel, PretrainedConfig
 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__()
-        # Load model and tokenizer
+class NewsClassifierConfig(PretrainedConfig):
+    model_type = "news_classifier"
+    
+    def __init__(self, max_length=128, **kwargs):
+        self.max_length = max_length
+        super().__init__(**kwargs)
+
+class NewsClassifier(PreTrainedModel):
+    config_class = NewsClassifierConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
         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):
-        """Tokenizes and pads the input text."""
+            
+    def forward(self, inputs):
         sequences = self.tokenizer.texts_to_sequences([inputs])
-        padded = pad_sequences(sequences, maxlen=128)
-        return padded
-
-    def _forward(self, inputs):
-        """Runs the model prediction."""
-        processed = self.preprocess(inputs)
-        predictions = self.model.predict(processed)
+        padded = pad_sequences(sequences, maxlen=self.config.max_length)
+        predictions = self.model.predict(padded)
         scores = predictions[0]
         label = "foxnews" if scores[0] > 0.5 else "nbc"
-        return [{"label": label, "score": float(scores[0] if label == "foxnews" else 1 - scores[0])}]
-
-    def postprocess(self, model_outputs):
-        """Returns the processed output."""
-        return model_outputs
+        return {"label": label, "score": float(scores[0] if label == "foxnews" else 1 - scores[0])}