Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
+
 os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' 
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 import logging
 import io
 import joblib
@@ -29,6 +30,7 @@ try:
 except (RuntimeError, ValueError) as e:
     logging.warning(f"Could not disable GPU for TensorFlow: {e}")
 
+
 # --- Model Loading ---
 def load_models():
     logging.info("Loading all models from the Hub...")
@@ -51,14 +53,14 @@ def load_models():
 
     try:
         xgb_model_path = hf_hub_download(
-            repo_id="muhalwan/california_housing_price_predictor", 
+            repo_id="muhalwan/california_housing_price_predictor",
             filename="xgb_model.joblib"
         )
         scaler_path = hf_hub_download(
-            repo_id="muhalwan/california_housing_price_predictor", 
+            repo_id="muhalwan/california_housing_price_predictor",
             filename="scaler.joblib"
         )
-        
+
         housing_model = joblib.load(xgb_model_path)
         housing_scaler = joblib.load(scaler_path)
         logging.info("Housing price model and scaler loaded successfully.")
@@ -68,14 +70,17 @@ def load_models():
 
     return tokenizer, sentiment_model, cat_dog_model, housing_model, housing_scaler
 
+
 # --- FastAPI App Initialization ---
 app = FastAPI()
 tokenizer, sentiment_model, cat_dog_model, housing_model, housing_scaler = load_models()
 app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
 
+
 class SentimentRequest(BaseModel):
     text: str
 
+
 class HousingRequest(BaseModel):
     MedInc: float
     HouseAge: float
@@ -86,11 +91,13 @@ class HousingRequest(BaseModel):
     Latitude: float
     Longitude: float
 
+
 # --- API Endpoints ---
 @app.get("/")
 async def read_root():
     return FileResponse('index.html')
 
+
 @app.post("/predict/sentiment")
 async def predict_sentiment(request: SentimentRequest):
     if not tokenizer or not sentiment_model:
@@ -101,13 +108,16 @@ async def predict_sentiment(request: SentimentRequest):
         with torch.no_grad():
             outputs = sentiment_model(**inputs)
         probabilities = F.softmax(outputs.logits, dim=-1).squeeze()
-        labels = ['Bearish', 'Bullish']
-        prediction = labels[torch.argmax(probabilities).item()]
+        id2label = sentiment_model.config.id2label
+        labels = [id2label[int(idx)] for idx in sorted([int(k) for k in id2label.keys()])]
+        pred_idx = torch.argmax(probabilities).item()
+        prediction = labels[pred_idx] if pred_idx < len(labels) else str(pred_idx)
         return {"prediction": prediction}
     except Exception as e:
         logging.error(f"Sentiment prediction error: {e}")
         raise HTTPException(status_code=500, detail="An error occurred during sentiment analysis.")
 
+
 @app.post("/predict/catdog")
 async def predict_catdog(file: UploadFile = File(...)):
     if not cat_dog_model:
@@ -124,7 +134,7 @@ async def predict_catdog(file: UploadFile = File(...)):
         img_array = np.expand_dims(img_array, axis=0)
         prob = cat_dog_model.predict(img_array, verbose=0)[0, 0]
         label = "Dog" if prob >= 0.5 else "Cat"
-        
+
         return {"prediction": label}
     except Exception as e:
         logging.error(f"Cat/Dog prediction error: {e}")