Update app.py

app.py

@@ -1,24 +1,93 @@
-import os
 from fastapi import FastAPI
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from transformers import pipeline
+import logging
 
-app = FastAPI()
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-MODEL_ID = "Adedayo2000/sentiment-model"
+app = FastAPI(
+    title="Sentiment Analysis API",
+    description="Analyze sentiment of text using BERT fine-tuned on IMDB dataset",
+    version="1.0.0",
+    docs_url="/docs",
+    redoc_url="/redoc"
+)
 
-hf_token = os.getenv("HF_TOKEN")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-print("Loading model...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=hf_token)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, token=hf_token)
-print("Model loaded successfully!")
+logger.info("Loading sentiment analysis model...")
+try:
+    classifier = pipeline("sentiment-analysis", model="Adedayo2000/sentiment-model")
+    logger.info("Model loaded successfully!")
+except Exception as e:
+    logger.error(f"Failed to load model: {e}")
+    classifier = None
+
+class TextInput(BaseModel):
+    text: str
+
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "text": "I love this product! It's amazing!"
+            }
+        }
 
 @app.get("/")
-def read_root():
-    return {"message": "API is live!", "model": MODEL_ID}
+def home():
+    return {
+        "message": "Sentiment Analysis API",
+        "status": "live",
+        "model": "Adedayo2000/sentiment-model",
+        "endpoints": {
+            "GET /": "API information",
+            "POST /predict": "Analyze sentiment of text",
+            "GET /docs": "Interactive API documentation (Swagger UI)",
+            "GET /redoc": "Alternative API documentation (ReDoc)",
+            "GET /health": "Health check"
+        }
+    }
+
+@app.get("/health")
+def health_check():
+    return {
+        "status": "healthy",
+        "model_loaded": classifier is not None
+    }
 
 @app.post("/predict")
-def predict(text: str):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    return {"logits": outputs.logits.tolist()}
+def predict(input: TextInput):
+    if not input.text or not input.text.strip():
+        return {
+            "error": "Text cannot be empty",
+            "example": {"text": "I love this product!"}
+        }
+    
+    if classifier is None:
+        return {
+            "error": "Model not loaded",
+            "status": "unavailable"
+        }
+    
+    try:
+        result = classifier(input.text)[0]
+        return {
+            "text": input.text,
+            "sentiment": result["label"],
+            "confidence": round(result["score"], 4)
+        }
+    except Exception as e:
+        logger.error(f"Prediction error: {e}")
+        return {
+            "error": "Prediction failed",
+            "details": str(e)
+        }
+