from fastapi import FastAPI
import lightgbm as lgb
from pydantic import BaseModel
from typing import List
from huggingface_hub import hf_hub_download

# Load model from Hugging Face
MODEL_REPO = "noisebop/test_model"
MODEL_FILENAME = "lightgbm_model.txt"

model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
model = lgb.Booster(model_file=model_path)

# Initialize FastAPI
app = FastAPI(title="ML Model API", version="1.0")

# Define request format
class InputData(BaseModel):
    features: List[float]  # List of numerical features

@app.get("/")
def home():
    return {"message": "Welcome to the ML API!"}

@app.post("/predict")
def predict(data: InputData):
    """Make a prediction using LightGBM"""
    try:
        prediction = model.predict([data.features])[0]
        return {"prediction": prediction}
    except Exception as e:
        return {"error": str(e)}