Add FastAPI inference API

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,14 @@
+FROM python:3.11.11-slim-bookworm
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,18 @@
+---
+title: Battery Analytics
+emoji: 🌍
+colorFrom: yellow
+colorTo: purple
+sdk: docker
+pinned: false
+license: mit
+---
+
+This Space serves the battery capacity prediction API.
+
+The service exposes:
+
+- `GET /health`
+- `POST /predict`
+
+The model weights and scalers are downloaded at runtime from the public model repo `Dharunkumar9/battery-capacity-predictor`, so the Space repo only needs code, not binary artifacts.

app.py

@@ -0,0 +1,43 @@
+from typing import Any, List
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+
+from predictor import get_predictor
+
+
+class PredictRequest(BaseModel):
+    battery_id: str = Field(default="B0005", description="Battery id or 0-based battery index")
+    window: List[List[float]] = Field(..., description="15 x 13 feature window")
+
+
+app = FastAPI(title="Battery Capacity Predictor API", version="1.0.0")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+predictor = get_predictor()
+
+
+@app.get("/")
+def root() -> dict[str, str]:
+    return {"status": "ok", "message": "Battery capacity prediction API is running"}
+
+
+@app.get("/health")
+def health() -> dict[str, str]:
+    return {"status": "healthy"}
+
+
+@app.post("/predict")
+def predict(request: PredictRequest) -> dict[str, Any]:
+    try:
+        return predictor.predict(request.window, battery_id=request.battery_id)
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail=str(exc)) from exc

predictor.py

@@ -0,0 +1,170 @@
+import json
+import os
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Dict, Mapping, Union
+
+import joblib
+import numpy as np
+import torch
+import torch.nn as nn
+from huggingface_hub import snapshot_download
+
+
+MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "Dharunkumar9/battery-capacity-predictor")
+BATTERY_ORDER = ["B0005", "B0006", "B0007", "B0018"]
+
+
+class PositionalEncoding(nn.Module):
+	def __init__(self, d_model: int, max_len: int = 500):
+		super().__init__()
+		pe = torch.zeros(max_len, d_model)
+		position = torch.arange(0, max_len).unsqueeze(1)
+		div_term = torch.exp(torch.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
+		pe[:, 0::2] = torch.sin(position * div_term)
+		pe[:, 1::2] = torch.cos(position * div_term)
+		self.register_buffer("pe", pe.unsqueeze(0))
+
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		return x + self.pe[:, : x.size(1), :]
+
+
+class BatteryTransformer(nn.Module):
+	def __init__(
+		self,
+		num_features: int,
+		d_model: int = 128,
+		nhead: int = 4,
+		num_layers: int = 2,
+		dim_feedforward: int = 256,
+		dropout: float = 0.1,
+		last_frac: float = 0.4,
+		last_weight: float = 3.0,
+	):
+		super().__init__()
+		self.input_proj = nn.Linear(num_features, d_model)
+		self.pos_encoder = PositionalEncoding(d_model)
+		encoder_layer = nn.TransformerEncoderLayer(
+			d_model=d_model,
+			nhead=nhead,
+			dim_feedforward=dim_feedforward,
+			dropout=dropout,
+			batch_first=True,
+		)
+		self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+		self.dropout = nn.Dropout(dropout)
+		self.regressor = nn.Linear(d_model, 1)
+		self.last_frac = last_frac
+		self.last_weight = last_weight
+
+	def forward(self, x: torch.Tensor) -> torch.Tensor:
+		seq_len = x.size(1)
+		x = self.input_proj(x)
+		x = self.pos_encoder(x)
+		x = self.encoder(x)
+
+		weights = torch.ones(seq_len, device=x.device)
+		last_start = int(seq_len * (1 - self.last_frac))
+		weights[last_start:] = self.last_weight
+		weights = weights / weights.sum()
+
+		x = (x * weights.unsqueeze(1)).sum(dim=1)
+		x = self.dropout(x)
+		return self.regressor(x).squeeze(-1)
+
+
+def _resolve_battery_id(battery_id: Union[str, int]) -> str:
+	if isinstance(battery_id, int):
+		if battery_id < 0 or battery_id >= len(BATTERY_ORDER):
+			raise ValueError(f"battery_id index must be between 0 and {len(BATTERY_ORDER) - 1}")
+		return BATTERY_ORDER[battery_id]
+
+	battery_id = str(battery_id).strip()
+	if battery_id not in BATTERY_ORDER:
+		raise ValueError(f"battery_id must be one of {BATTERY_ORDER} or a 0-based index")
+	return battery_id
+
+
+def _normalize_window(window: Any, expected_rows: int, expected_cols: int) -> np.ndarray:
+	array = np.asarray(window, dtype=np.float32)
+	if array.ndim == 1:
+		if array.size != expected_rows * expected_cols:
+			raise ValueError(f"window must contain {expected_rows * expected_cols} values")
+		array = array.reshape(expected_rows, expected_cols)
+
+	if array.shape != (expected_rows, expected_cols):
+		raise ValueError(f"window must have shape ({expected_rows}, {expected_cols})")
+
+	return array
+
+
+def _download_artifacts() -> Path:
+	return Path(
+		snapshot_download(
+			repo_id=MODEL_REPO_ID,
+			repo_type="model",
+			allow_patterns=["config.json", "pytorch_model.bin", "x_scalers.pkl", "y_scalers.pkl"],
+		)
+	)
+
+
+class BatteryPredictor:
+	def __init__(self) -> None:
+		artifact_dir = _download_artifacts()
+		config = json.loads((artifact_dir / "config.json").read_text())
+		self.window_size = int(config["window_size"])
+		self.num_features = int(config["num_features"])
+
+		self.x_scalers = joblib.load(artifact_dir / "x_scalers.pkl")
+		self.y_scalers = joblib.load(artifact_dir / "y_scalers.pkl")
+
+		self.model = BatteryTransformer(
+			num_features=self.num_features,
+			d_model=int(config["d_model"]),
+			nhead=int(config["nhead"]),
+			num_layers=int(config["num_layers"]),
+			dim_feedforward=int(config["dim_feedforward"]),
+			dropout=float(config["dropout"]),
+		).to("cpu")
+
+		state_dict = torch.load(artifact_dir / "pytorch_model.bin", map_location="cpu")
+		self.model.load_state_dict(state_dict)
+		self.model.eval()
+
+	def predict(self, window: Any, battery_id: Union[str, int] = "B0005") -> Dict[str, Any]:
+		battery_key = _resolve_battery_id(battery_id)
+		window_array = _normalize_window(window, self.window_size, self.num_features)
+
+		x_scaler = self.x_scalers[battery_key]
+		y_scaler = self.y_scalers[battery_key]
+
+		scaled_window = x_scaler.transform(window_array)
+		tensor = torch.tensor(scaled_window[None, :, :], dtype=torch.float32)
+
+		with torch.no_grad():
+			scaled_prediction = float(self.model(tensor).item())
+
+		predicted_capacity = float(y_scaler.inverse_transform([[scaled_prediction]])[0, 0])
+
+		return {
+			"battery_id": battery_key,
+			"window_size": self.window_size,
+			"num_features": self.num_features,
+			"predicted_capacity": predicted_capacity,
+			"scaled_prediction": scaled_prediction,
+		}
+
+
+@lru_cache(maxsize=1)
+def get_predictor() -> BatteryPredictor:
+	return BatteryPredictor()
+
+
+def predict_from_request(payload: Mapping[str, Any]) -> Dict[str, Any]:
+	if not isinstance(payload, Mapping):
+		raise TypeError("payload must be a mapping with battery_id and window")
+
+	if "window" not in payload:
+		raise ValueError("payload must include a window field")
+
+	return get_predictor().predict(payload["window"], battery_id=payload.get("battery_id", "B0005"))

requirements.txt

@@ -0,0 +1,7 @@
+fastapi>=0.110.0
+uvicorn[standard]>=0.29.0
+torch>=2.2.0
+numpy>=1.24.0
+joblib>=1.3.0
+scikit-learn==1.7.0
+huggingface_hub>=0.23.0