Spaces:
Running on Zero
Running on Zero
| import os | |
| import joblib | |
| from fastapi import FastAPI, Query | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from huggingface_hub import hf_hub_download | |
| import pandas as pd | |
| import numpy as np | |
| # Local paths (used when these files exist, e.g. during local development) | |
| model_uri = "../../../ml/models/FD004GBM_BAG" | |
| data_uri = "../../../data/FE_data/test_FD004" | |
| # Hugging Face Hub locations (used when the local files are not found) | |
| HF_MODEL_REPO = "ZephyrNTempest/JetBoost-GBM" | |
| HF_MODEL_FILE = "models/FD004GBM_BAG" | |
| HF_DATA_REPO = "ZephyrNTempest/C-MAPSS" | |
| HF_DATA_FILE = "data/FE_data/test_FD004" | |
| HF_DATA_REPO_TYPE = ( | |
| "dataset" # "dataset" or "model", whichever you uploaded the data to | |
| ) | |
| def resolve_path(local_path, repo_id, filename, repo_type="model"): | |
| """Return a local file path, downloading from the Hugging Face Hub if needed. | |
| Args: | |
| local_path (str): Local path to use if the file exists. | |
| repo_id (str): Hub repo id, e.g. "username/repo". | |
| filename (str): File name inside the Hub repo. | |
| repo_type (str): "model" or "dataset". | |
| Returns: | |
| str: Path to a readable local copy of the file. | |
| """ | |
| if os.path.exists(local_path): | |
| return local_path | |
| return hf_hub_download(repo_id=repo_id, filename=filename, repo_type=repo_type) | |
| # Load the model | |
| models = joblib.load(resolve_path(model_uri, HF_MODEL_REPO, HF_MODEL_FILE)) | |
| app = FastAPI() | |
| # Allow the GitHub Pages site (and local dev) to call this API. | |
| # Set ALLOWED_ORIGINS as a comma-separated list in the Space's variables. | |
| allowed_origins = os.getenv( | |
| "ALLOWED_ORIGINS", "http://localhost:5173,https://YOUR_USERNAME.github.io" | |
| ).split(",") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=[o.strip() for o in allowed_origins], | |
| allow_methods=["GET"], | |
| allow_headers=["*"], | |
| ) | |
| # Prune test set features | |
| def prune_features(test, drop_pc2=False): | |
| """Drop min/last/max/delta feature columns (and pc2 if asked) from train and test.""" | |
| patterns = ["min", "last", "max", "delta"] | |
| if drop_pc2: | |
| patterns.append("pc2") | |
| to_drop = [c for c in test.columns if any(p in c for p in patterns)] | |
| return test.drop(columns=to_drop, errors="ignore") | |
| df = pd.read_parquet( | |
| resolve_path(data_uri, HF_DATA_REPO, HF_DATA_FILE, HF_DATA_REPO_TYPE) | |
| ) | |
| df = df.drop(columns="RUL") | |
| df = df.set_index("Engine") | |
| df = prune_features(df) | |
| # Average the preditions of the bagged models | |
| class Bag: | |
| def __init__(self, models): | |
| self.models = models | |
| def predict(self, X): | |
| return np.mean([m.predict(X) for m in self.models], axis=0) | |
| bag = Bag(models) | |
| def inference(target_idxs: list[int] = Query(...), cycles: list[int] = Query(...)): | |
| """Return FD004GBM model predictions for a batch of engines, | |
| each at its own time cycle. | |
| Args: | |
| target_idxs (list[int]): Engine indices (e.g. 1..100) to predict on. | |
| cycles (list[int]): Time cycle for each engine. cycles[i] belongs to | |
| target_idxs[i], so both lists must be the same length. | |
| Returns: | |
| dict: {"Predicted RUL": [...]} with one prediction per engine, | |
| in the same order as target_idxs. | |
| """ | |
| lookup = dict(zip(target_idxs, cycles)) | |
| wanted = df.index.map(lookup) | |
| targets = df[df["Time"] == wanted] | |
| targets = targets.loc[target_idxs] | |
| rul_pred = bag.predict(targets) | |
| return {"Predicted RUL": rul_pred.tolist()} | |