Create pipeline.py

pipeline.py

@@ -0,0 +1,64 @@
+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 NewsClassifierConfig(PretrainedConfig):
+    model_type = "news_classifier"
+    
+    def __init__(
+        self,
+        max_length=41,  # Modified to match model input shape
+        vocab_size=74934,  # Modified based on embedding layer size
+        embedding_dim=128,  # Added to match model architecture
+        hidden_size=64,  # Matches final LSTM layer
+        num_labels=2,
+        **kwargs
+    ):
+        self.max_length = max_length
+        self.vocab_size = vocab_size
+        self.embedding_dim = embedding_dim
+        self.hidden_size = hidden_size
+        self.num_labels = num_labels
+        super().__init__(**kwargs)
+
+class NewsClassifier(PreTrainedModel):
+    config_class = NewsClassifierConfig
+    base_model_prefix = "news_classifier"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = None  # Will be loaded in post_init
+        self.tokenizer = None
+        
+    def post_init(self):
+        """Load model and tokenizer after initialization"""
+        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 forward(self, text_input):
+        if not self.model or not self.tokenizer:
+            self.post_init()
+            
+        if isinstance(text_input, str):
+            text_input = [text_input]
+            
+        sequences = self.tokenizer.texts_to_sequences(text_input)
+        padded = pad_sequences(sequences, maxlen=self.config.max_length)
+        predictions = self.model.predict(padded, verbose=0)
+        
+        results = []
+        for pred in predictions:
+            # Convert from 2-class output to single score
+            score = float(pred[1])  # Assuming pred[1] is the probability for "foxnews"
+            label = "foxnews" if score > 0.5 else "nbc"
+            results.append({
+                "label": label,
+                "score": score if label == "foxnews" else 1 - score
+            })
+        
+        return results[0] if len(text_input) == 1 else results