Attempt to resolve deployment issue

README.md

@@ -1,4 +1,12 @@
 ---
+title: News Source Classifier
+emoji: 📰
+colorFrom: blue
+colorTo: red
+sdk: fastapi
+sdk_version: 0.95.2
+app_file: app.py
+pinned: false
 language: en
 license: mit
 tags:
@@ -43,40 +51,14 @@ This model classifies news headlines as either Fox News or NBC News using an LST
 
 ## Usage
 
-You can use this model directly with a FastAPI endpoint:
+You can use this model through the FastAPI endpoint:
 
 ```python
 import requests
 
 # Make a prediction
 response = requests.post(
-    "https://huggingface.co/Jiahuita/NewsSourceClassification/predict",
+    "https://huggingface.co/Jiahuita/NewsSourceClassification",
     json={"text": "Your news headline here"}
 )
-print(response.json())
-```
-
-Or use it locally:
-
-```python
-from transformers import pipeline
-
-classifier = pipeline("text-classification", model="Jiahuita/NewsSourceClassification")
-result = classifier("Your news headline here")
-print(result)
-```
-
-## Limitations and Bias
-
-This model has been trained on news headlines from specific sources and time periods, which may introduce certain biases. Users should be aware of these limitations when using the model.
-
-## Training
-
-The model was trained using:
-- TensorFlow 2.13.0
-- LSTM architecture
-- Binary cross-entropy loss
-- Adam optimizer
-
-## License
-This project is licensed under the MIT License.
+print(response.json())

app.py

@@ -0,0 +1,15 @@
+from transformers import pipeline
+from fastapi import FastAPI
+from pydantic import BaseModel
+
+app = FastAPI()
+
+class TextInput(BaseModel):
+    text: str
+
+classifier = pipeline("text-classification", model="./")
+
+@app.post("/predict")
+async def predict(input_data: TextInput):
+    result = classifier(input_data.text)
+    return result

pipeline.py

@@ -1,52 +1,35 @@
-from transformers import PreTrainedModel, PretrainedConfig
-from tensorflow.keras.models import load_model
-from tensorflow.keras.preprocessing.text import tokenizer_from_json
+from transformers import Pipeline
+import tensorflow as tf
 from tensorflow.keras.preprocessing.sequence import pad_sequences
-import os
-import numpy as np
 import json
+import os
 
-class NewsClassifierConfig(PretrainedConfig):
-    model_type = "news_classifier"
-    
-    def __init__(self, max_length=128, **kwargs):
-        self.max_length = max_length
-        super().__init__(**kwargs)
+def load_tokenizer(tokenizer_path):
+    with open(tokenizer_path, 'r') as f:
+        return json.load(f)
 
-class NewsClassifier(PreTrainedModel):
-    config_class = NewsClassifierConfig
-    
-    def __init__(self, config):
-        super().__init__(config)
+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)
         
-        self.model = load_model(model_path)
-        with open(tokenizer_path, 'r') as f:
-            tokenizer_data = json.load(f)
-            self.tokenizer = tokenizer_from_json(tokenizer_data)
+    def __call__(self, texts, **kwargs):
+        if isinstance(texts, str):
+            texts = [texts]
 
-    def forward(self, text_input):
-        if isinstance(text_input, str):
-            sequences = self.tokenizer.texts_to_sequences([text_input])
-        else:
-            sequences = self.tokenizer.texts_to_sequences(text_input)
-            
-        padded = pad_sequences(sequences, maxlen=self.config.max_length)
+        sequences = self.tokenizer.texts_to_sequences(texts)
+        padded = pad_sequences(sequences, maxlen=128)
+        
         predictions = self.model.predict(padded)
         
         results = []
-        for score in predictions:
-            label = "foxnews" if score[0] > 0.5 else "nbc"
-            results.append({
-                "label": label,
-                "score": float(score[0] if label == "foxnews" else 1 - score[0])
-            })
-        
-        return results[0] if isinstance(text_input, str) else results
-
-    @classmethod
-    def from_pretrained(cls, model_path, **kwargs):
-        config = NewsClassifierConfig.from_pretrained(model_path)
-        model = cls(config)
-        return model
+        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

requirements.txt

@@ -2,3 +2,6 @@ tensorflow>=2.10.0
 transformers>=4.46.3
 numpy>=1.19.2
 scikit-learn>=0.24.2
+fastapi>=0.68.0
+uvicorn>=0.15.0
+pydantic>=1.8.2