app.py

# app.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
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
from typing import Union, List

app = FastAPI()

# Load model and tokenizer
model = load_model('news_classifier.h5')
with open('tokenizer.json', 'r') as f:
    tokenizer_data = json.load(f)
    tokenizer = tokenizer_from_json(tokenizer_data)

class PredictionInput(BaseModel):
    text: Union[str, List[str]]

class PredictionOutput(BaseModel):
    label: str
    score: float

@app.post("/predict")
async def predict(input_data: PredictionInput):
    try:
        # Convert input to list if it's a single string
        texts = input_data.text if isinstance(input_data.text, list) else [input_data.text]
        
        # Preprocess
        sequences = tokenizer.texts_to_sequences(texts)
        padded = pad_sequences(sequences, maxlen=41)  # Use your model's expected input length
        
        # Predict
        predictions = model.predict(padded)
        
        # Format results
        results = []
        for pred in predictions:
            score = float(pred[1])  # Assuming binary classification
            label = "foxnews" if score > 0.5 else "nbc"
            results.append({
                "label": label,
                "score": score if label == "foxnews" else 1 - score
            })
        
        # Return single result if input was single string
        return results[0] if isinstance(input_data.text, str) else results
            
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"message": "News Classifier API is running"}