app/main.py

# Hugging Face Space: app/main.py
import joblib
import pandas as pd
import numpy as np
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from huggingface_hub import hf_hub_download
from typing import List

# --- Pydantic Models (Input and Output validation) ---
class PredictionFeatures(BaseModel):
    Field: str
    GPA: float
    Leadership_Positions: int
    Research_Experience: int
    Industry_Certifications: int
    Extracurricular_Activities: int
    Internships: int
    Projects: int
    Field_Specific_Courses: int
    Coding_Skills: int
    Communication_Skills: int
    Problem_Solving_Skills: int
    Teamwork_Skills: int
    Analytical_Skills: int
    Presentation_Skills: int
    Networking_Skills: int

class PredictionInput(BaseModel):
    features: PredictionFeatures

class CareerPrediction(BaseModel):
    career: str
    probability: float

class PredictionOutput(BaseModel):
    top_predictions: List[CareerPrediction]

# --- FastAPI App ---
app = FastAPI()

# --- Load Models from your Hugging Face Hub Model Repository ---
MODEL = None
LABEL_ENCODER = None
FEATURE_COLUMNS = None

try:
    repo_id = "ash47/margdarshak-career-predictor"
    # Define a local, writable cache directory
    CACHE_DIR = "./model_cache"
    print(f"--- Loading models from Hub repo: {repo_id} to cache: {CACHE_DIR} ---")

    # FINAL FIX: Add the cache_dir parameter to each download call
    model_path = hf_hub_download(
        repo_id=repo_id,
        filename="career_prediction_model_pipeline.joblib",
        cache_dir=CACHE_DIR
    )
    MODEL = joblib.load(model_path)
    print("-> Model pipeline loaded.")

    label_encoder_path = hf_hub_download(
        repo_id=repo_id,
        filename="career_label_encoder.joblib",
        cache_dir=CACHE_DIR
    )
    LABEL_ENCODER = joblib.load(label_encoder_path)
    print("-> Label encoder loaded.")

    feature_columns_path = hf_hub_download(
        repo_id=repo_id,
        filename="career_feature_columns.joblib",
        cache_dir=CACHE_DIR
    )
    FEATURE_COLUMNS = joblib.load(feature_columns_path)
    print("-> Feature columns loaded...")
    
    print("--- All models loaded successfully. ---")

except Exception as e:
    print(f"FATAL: Error loading models from Hugging Face Hub: {e}")


@app.get("/")
def read_root():
    return {"status": "ok", "models_loaded": all([MODEL, LABEL_ENCODER, FEATURE_COLUMNS])}

@app.post("/predict", response_model=PredictionOutput)
def predict(payload: PredictionInput):
    if not all([MODEL, LABEL_ENCODER, FEATURE_COLUMNS]):
        raise HTTPException(status_code=503, detail="Models are not available or failed to load.")

    try:
        input_df = pd.DataFrame([payload.features.dict()])
        input_df = input_df.reindex(columns=FEATURE_COLUMNS, fill_value=0)
        probabilities = MODEL.predict_proba(input_df)[0]
        top_n_indices = np.argsort(probabilities)[::-1][:5]
        top_predictions = [
            CareerPrediction(
                career=LABEL_ENCODER.classes_[i],
                probability=round(float(probabilities[i]), 4)
            ) for i in top_n_indices
        ]
        return PredictionOutput(top_predictions=top_predictions)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")