Modified app and readme

README.md

@@ -56,8 +56,43 @@ You can use this model directly with a FastAPI endpoint:
 ```python
 import requests
 
+# Make a prediction
 response = requests.post(
-    "https://huggingface.co/Jiahuita/NewsSourceClassification",
+    "https://huggingface.co/Jiahuita/NewsSourceClassification/predict",
     json={"text": "Your news headline here"}
 )
-print(response.json())
+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)
+```
+
+Example response:
+```json
+{
+    "label": "foxnews",
+    "score": 0.875
+}
+```
+
+## 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.

app.py

@@ -1,15 +1,84 @@
-from transformers import pipeline
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-
-app = FastAPI()
+from transformers import Pipeline
+import tensorflow as tf
+from tensorflow.keras.preprocessing.sequence import pad_sequences
+import json
+import os
 
 class TextInput(BaseModel):
     text: str
 
-classifier = pipeline("text-classification", model="./")
+app = FastAPI(
+    title="News Source Classifier",
+    description="A model to classify news headlines as either Fox News or NBC News",
+    version="1.0.0"
+)
+
+class NewsClassificationPipeline(Pipeline):
+    def __init__(self):
+        super().__init__()
+        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')
+        with open(tokenizer_path, 'r') as f:
+            tokenizer_data = json.load(f)
+            self.tokenizer = tf.keras.preprocessing.text.tokenizer_from_json(tokenizer_data)
+    
+    def __call__(self, text):
+        if isinstance(text, str):
+            text = [text]
+
+        sequences = self.tokenizer.texts_to_sequences(text)
+        padded = pad_sequences(sequences, maxlen=128)
+        
+        predictions = self.model.predict(padded)
+
+        results = []
+        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 len(results) == 1 else results
+
+try:
+    classifier = NewsClassificationPipeline()
+except Exception as e:
+    print(f"Error initializing model: {str(e)}")
+    raise
+
+@app.get("/")
+async def root():
+    return {
+        "message": "News Source Classification API",
+        "usage": "Send POST request to /predict with {'text': 'your news headline'}"
+    }
 
 @app.post("/predict")
 async def predict(input_data: TextInput):
-    result = classifier(input_data.text)
-    return result
+    try:
+        result = classifier(input_data.text)
+        return result
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/examples")
+async def examples():
+    return {
+        "examples": [
+            {
+                "title": "Crime News Headline",
+                "text": "Wife of murdered Minnesota pastor hired 3 men to kill husband after affair: police"
+            },
+            {
+                "title": "Science News Headline",
+                "text": "Scientists discover breakthrough in renewable energy research"
+            },
+            {
+                "title": "Political News Headline",
+                "text": "Presidential candidates face off in heated debate over climate policies"
+            }
+        ]
+    }