from fastapi import FastAPI
from pydantic import BaseModel
import numpy as np
import pickle
from huggingface_hub import hf_hub_download

app = FastAPI()

# Download your model pickle from the Hub on startup
model_path = hf_hub_download(repo_id="nataliegref/test-model", filename="trained_causal_model_v2.pkl")

with open(model_path, "rb") as f:
    model = pickle.load(f)

class InputData(BaseModel):
    X: list  

@app.post("/predict")
def predict(data: InputData):
    X = np.array(data.X)
    effect = model.effect(X).tolist()
    return {"effect": effect}