Add remote TimesFM service integration
This commit is contained in:
2569718930@qq.com
+6
-1
@@ -36,12 +36,17 @@ POLYWEATHER_BOT_GROUP_INVITE_URL=
|
||||
OPEN_METEO_CACHE_TTL_SEC=7200
|
||||
OPEN_METEO_ENSEMBLE_CACHE_TTL_SEC=7200
|
||||
OPEN_METEO_MULTI_MODEL_CACHE_TTL_SEC=7200
|
||||
OPEN_METEO_MULTI_MODEL_CACHE_VERSION=v2
|
||||
OPEN_METEO_MULTI_MODEL_CACHE_VERSION=v4
|
||||
OPEN_METEO_RATE_LIMIT_COOLDOWN_SEC=900
|
||||
OPEN_METEO_RATE_CACHE_TTL_SEC=3600
|
||||
OPEN_METEO_MIN_CALL_INTERVAL_SEC=3
|
||||
METAR_CACHE_TTL_SEC=600
|
||||
METEOBLUE_CACHE_TTL_SEC=7200
|
||||
POLYWEATHER_TIMESFM_ENABLED=false
|
||||
POLYWEATHER_TIMESFM_SERVICE_URL=http://polyweather_timesfm:8011
|
||||
POLYWEATHER_TIMESFM_TIMEOUT_SEC=12
|
||||
POLYWEATHER_TIMESFM_HISTORY_LIMIT=60
|
||||
POLYWEATHER_TIMESFM_MIN_HISTORY_POINTS=14
|
||||
|
||||
########################################
|
||||
# 4) Auth / entitlement
|
||||
|
||||
@@ -162,6 +162,21 @@ docker compose logs -f polyweather | egrep "polymarket wallet activity watcher s
|
||||
| `/diag` | Startup diagnostics |
|
||||
| `/help` | Help and usage |
|
||||
|
||||
## Optional TimesFM Service
|
||||
|
||||
PolyWeather can consume a separate official TimesFM inference service for 1 to 3 day forecasts.
|
||||
|
||||
1. Enable it in `.env`:
|
||||
- `POLYWEATHER_TIMESFM_ENABLED=true`
|
||||
- `POLYWEATHER_TIMESFM_SERVICE_URL=http://polyweather_timesfm:8011`
|
||||
2. Start the optional service profile:
|
||||
|
||||
```bash
|
||||
docker compose --profile timesfm up --build polyweather_timesfm
|
||||
```
|
||||
|
||||
Service implementation lives in [`timesfm_service/`](timesfm_service/).
|
||||
|
||||
## Documentation Index
|
||||
|
||||
- Chinese overview: [README_ZH.md](README_ZH.md)
|
||||
|
||||
@@ -31,3 +31,15 @@ services:
|
||||
- "8000:8000"
|
||||
# UID/GID are mainly useful on Linux hosts to avoid root-owned output files.
|
||||
user: "${UID:-1000}:${GID:-1000}"
|
||||
|
||||
polyweather_timesfm:
|
||||
build:
|
||||
context: .
|
||||
dockerfile: timesfm_service/Dockerfile
|
||||
container_name: polyweather_timesfm
|
||||
restart: unless-stopped
|
||||
profiles:
|
||||
- timesfm
|
||||
environment:
|
||||
TIMESFM_MODEL_ID: ${TIMESFM_MODEL_ID:-google/timesfm-2.5-200m-pytorch}
|
||||
TIMESFM_PORT: ${POLYWEATHER_TIMESFM_PORT:-8011}
|
||||
|
||||
@@ -186,6 +186,7 @@ export interface ModelComparison {
|
||||
ICON?: number;
|
||||
GEM?: number;
|
||||
JMA?: number;
|
||||
TimesFM?: number;
|
||||
MGM?: number;
|
||||
NWS?: number;
|
||||
}
|
||||
|
||||
@@ -559,6 +559,38 @@ class NwsOpenMeteoSourceMixin:
|
||||
record_source_call("open_meteo", "multi_model", "empty", (time.perf_counter() - started) * 1000.0)
|
||||
return None
|
||||
|
||||
timesfm_meta = {}
|
||||
try:
|
||||
from src.models.timesfm_adapter import (
|
||||
TIMESFM_MODEL_NAME,
|
||||
predict_timesfm_daily,
|
||||
)
|
||||
|
||||
timesfm_result = predict_timesfm_daily(
|
||||
city_name=city,
|
||||
forecast_dates=dates,
|
||||
daily_model_forecasts=daily_forecasts,
|
||||
)
|
||||
timesfm_predictions = (
|
||||
timesfm_result.get("predictions", {})
|
||||
if isinstance(timesfm_result, dict)
|
||||
else {}
|
||||
)
|
||||
if isinstance(timesfm_predictions, dict):
|
||||
for date_str, raw_value in timesfm_predictions.items():
|
||||
parsed = round(float(raw_value), 1) if raw_value is not None else None
|
||||
if parsed is None:
|
||||
continue
|
||||
daily_forecasts.setdefault(date_str, {})[TIMESFM_MODEL_NAME] = parsed
|
||||
if isinstance(timesfm_result, dict):
|
||||
timesfm_meta = {
|
||||
key: value
|
||||
for key, value in timesfm_result.items()
|
||||
if key != "predictions"
|
||||
}
|
||||
except Exception as exc:
|
||||
logger.warning(f"TimesFM adapter failed for {city}: {exc}")
|
||||
|
||||
# 今天的预报 (向后兼容)
|
||||
today_date = dates[0] if dates else None
|
||||
forecasts = daily_forecasts.get(today_date, {})
|
||||
@@ -574,6 +606,7 @@ class NwsOpenMeteoSourceMixin:
|
||||
"daily_forecasts": daily_forecasts, # 按天 {"2026-02-23": {...}, "2026-02-24": {...}}
|
||||
"dates": dates,
|
||||
"unit": "fahrenheit" if use_fahrenheit else "celsius",
|
||||
"timesfm_meta": timesfm_meta,
|
||||
}
|
||||
with self._multi_model_cache_lock:
|
||||
self._multi_model_cache[cache_key] = {
|
||||
|
||||
@@ -111,8 +111,8 @@ class WeatherDataCollector(OpenMeteoCacheMixin, SettlementSourceMixin, MetarSour
|
||||
os.getenv("OPEN_METEO_MULTI_MODEL_CACHE_TTL_SEC", "900")
|
||||
)
|
||||
self.multi_model_cache_version = str(
|
||||
os.getenv("OPEN_METEO_MULTI_MODEL_CACHE_VERSION", "v2")
|
||||
).strip() or "v2"
|
||||
os.getenv("OPEN_METEO_MULTI_MODEL_CACHE_VERSION", "v4")
|
||||
).strip() or "v4"
|
||||
self._open_meteo_cache: Dict[str, Dict] = {}
|
||||
self._ensemble_cache: Dict[str, Dict] = {}
|
||||
self._multi_model_cache: Dict[str, Dict] = {}
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import requests
|
||||
from loguru import logger
|
||||
|
||||
from src.utils.metrics import record_source_call
|
||||
|
||||
|
||||
TIMESFM_MODEL_NAME = "TimesFM"
|
||||
TIMESFM_DEFAULT_MODEL_ID = "google/timesfm-2.5-200m-pytorch"
|
||||
|
||||
|
||||
def _sf(value: object) -> Optional[float]:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
return float(value)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _parse_date(raw: object) -> Optional[datetime]:
|
||||
text = str(raw or "").strip()
|
||||
if not text:
|
||||
return None
|
||||
try:
|
||||
return datetime.strptime(text, "%Y-%m-%d")
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _get_default_history_file() -> str:
|
||||
project_root = os.path.dirname(
|
||||
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
)
|
||||
return os.path.join(project_root, "data", "daily_records.json")
|
||||
|
||||
|
||||
def _is_enabled() -> bool:
|
||||
return str(os.getenv("POLYWEATHER_TIMESFM_ENABLED", "false")).strip().lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
}
|
||||
|
||||
|
||||
def _get_service_url() -> str:
|
||||
return str(os.getenv("POLYWEATHER_TIMESFM_SERVICE_URL", "")).strip().rstrip("/")
|
||||
|
||||
|
||||
def _load_actual_history(
|
||||
city_name: str,
|
||||
history_file: Optional[str] = None,
|
||||
max_points: Optional[int] = None,
|
||||
) -> List[Dict[str, object]]:
|
||||
from src.analysis.deb_algorithm import load_history
|
||||
from src.data_collection.city_registry import ALIASES
|
||||
|
||||
city_key = ALIASES.get(
|
||||
str(city_name or "").strip().lower(),
|
||||
str(city_name or "").strip().lower(),
|
||||
)
|
||||
if not city_key:
|
||||
return []
|
||||
|
||||
data = load_history(history_file or _get_default_history_file())
|
||||
city_data = data.get(city_key) if isinstance(data, dict) else None
|
||||
if not isinstance(city_data, dict):
|
||||
return []
|
||||
|
||||
rows: List[Dict[str, object]] = []
|
||||
for date_key, record in city_data.items():
|
||||
if not isinstance(record, dict):
|
||||
continue
|
||||
stamp = _parse_date(date_key)
|
||||
actual = _sf(record.get("actual_high"))
|
||||
if stamp is None or actual is None:
|
||||
continue
|
||||
rows.append(
|
||||
{
|
||||
"timestamp": stamp.strftime("%Y-%m-%d"),
|
||||
"value": round(actual, 1),
|
||||
}
|
||||
)
|
||||
|
||||
rows.sort(key=lambda item: str(item.get("timestamp") or ""))
|
||||
if max_points and max_points > 0:
|
||||
rows = rows[-max_points:]
|
||||
return rows
|
||||
|
||||
|
||||
def predict_timesfm_daily(
|
||||
*,
|
||||
city_name: str,
|
||||
forecast_dates: List[str],
|
||||
daily_model_forecasts: Optional[Dict[str, Dict[str, object]]] = None,
|
||||
history_file: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
started = time.perf_counter()
|
||||
service_url = _get_service_url()
|
||||
timeout_sec = float(os.getenv("POLYWEATHER_TIMESFM_TIMEOUT_SEC", "12"))
|
||||
max_history_points = int(os.getenv("POLYWEATHER_TIMESFM_HISTORY_LIMIT", "60"))
|
||||
min_history_points = int(os.getenv("POLYWEATHER_TIMESFM_MIN_HISTORY_POINTS", "14"))
|
||||
|
||||
if not _is_enabled():
|
||||
return {
|
||||
"predictions": {},
|
||||
"enabled": False,
|
||||
"reason": "disabled",
|
||||
}
|
||||
|
||||
if not service_url:
|
||||
return {
|
||||
"predictions": {},
|
||||
"enabled": False,
|
||||
"reason": "service_not_configured",
|
||||
}
|
||||
|
||||
normalized_dates = [
|
||||
str(date_str or "").strip()
|
||||
for date_str in (forecast_dates or [])
|
||||
if _parse_date(date_str) is not None
|
||||
]
|
||||
normalized_dates = list(dict.fromkeys(normalized_dates))
|
||||
if not normalized_dates:
|
||||
return {
|
||||
"predictions": {},
|
||||
"enabled": True,
|
||||
"reason": "no_valid_future_dates",
|
||||
}
|
||||
|
||||
history_rows = _load_actual_history(
|
||||
city_name=city_name,
|
||||
history_file=history_file,
|
||||
max_points=max_history_points,
|
||||
)
|
||||
if len(history_rows) < min_history_points:
|
||||
return {
|
||||
"predictions": {},
|
||||
"enabled": True,
|
||||
"reason": "insufficient_history",
|
||||
"history_count": len(history_rows),
|
||||
}
|
||||
|
||||
payload = {
|
||||
"city": city_name,
|
||||
"series_frequency": "D",
|
||||
"series_kind": "actual_high",
|
||||
"series": history_rows,
|
||||
"future_dates": normalized_dates,
|
||||
"daily_model_forecasts": daily_model_forecasts or {},
|
||||
}
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{service_url}/predict/daily",
|
||||
json=payload,
|
||||
timeout=timeout_sec,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
raw_predictions = data.get("predictions", {}) if isinstance(data, dict) else {}
|
||||
|
||||
predictions: Dict[str, float] = {}
|
||||
for date_str in normalized_dates:
|
||||
parsed = _sf((raw_predictions or {}).get(date_str))
|
||||
if parsed is None:
|
||||
continue
|
||||
predictions[date_str] = round(parsed, 1)
|
||||
|
||||
record_source_call(
|
||||
"timesfm",
|
||||
"predict",
|
||||
"success" if predictions else "empty",
|
||||
(time.perf_counter() - started) * 1000.0,
|
||||
)
|
||||
return {
|
||||
"predictions": predictions,
|
||||
"enabled": True,
|
||||
"reason": "ok" if predictions else "empty",
|
||||
"history_count": len(history_rows),
|
||||
"service_url": service_url,
|
||||
"model": data.get("model") if isinstance(data, dict) else None,
|
||||
"model_id": data.get("model_id") if isinstance(data, dict) else None,
|
||||
"series_frequency": "D",
|
||||
"series_kind": "actual_high",
|
||||
"quantiles": data.get("quantiles") if isinstance(data, dict) else None,
|
||||
}
|
||||
except Exception as exc:
|
||||
logger.warning(f"TimesFM remote request failed for {city_name}: {exc}")
|
||||
record_source_call(
|
||||
"timesfm",
|
||||
"predict",
|
||||
"error",
|
||||
(time.perf_counter() - started) * 1000.0,
|
||||
)
|
||||
return {
|
||||
"predictions": {},
|
||||
"enabled": True,
|
||||
"reason": "request_failed",
|
||||
"history_count": len(history_rows),
|
||||
"service_url": service_url,
|
||||
"error": str(exc),
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from src.models.timesfm_adapter import predict_timesfm_daily
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload):
|
||||
self._payload = payload
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
def test_predict_timesfm_daily_disabled_by_default(monkeypatch):
|
||||
monkeypatch.delenv("POLYWEATHER_TIMESFM_ENABLED", raising=False)
|
||||
monkeypatch.delenv("POLYWEATHER_TIMESFM_SERVICE_URL", raising=False)
|
||||
|
||||
result = predict_timesfm_daily(
|
||||
city_name="ankara",
|
||||
forecast_dates=["2026-03-15", "2026-03-16"],
|
||||
)
|
||||
|
||||
assert result["predictions"] == {}
|
||||
assert result["enabled"] is False
|
||||
assert result["reason"] == "disabled"
|
||||
|
||||
|
||||
def test_predict_timesfm_daily_posts_actual_history_to_remote_service(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
def fake_load_history(_filepath):
|
||||
return {
|
||||
"ankara": {
|
||||
"2026-03-01": {"actual_high": 10.0},
|
||||
"2026-03-02": {"actual_high": 11.0},
|
||||
"2026-03-03": {"actual_high": 12.0},
|
||||
"2026-03-04": {"actual_high": 13.0},
|
||||
"2026-03-05": {"actual_high": 14.0},
|
||||
"2026-03-06": {"actual_high": 15.0},
|
||||
"2026-03-07": {"actual_high": 16.0},
|
||||
"2026-03-08": {"actual_high": 15.0},
|
||||
"2026-03-09": {"actual_high": 14.0},
|
||||
"2026-03-10": {"actual_high": 15.0},
|
||||
"2026-03-11": {"actual_high": 16.0},
|
||||
"2026-03-12": {"actual_high": 17.0},
|
||||
"2026-03-13": {"actual_high": 16.0},
|
||||
"2026-03-14": {"actual_high": 15.0},
|
||||
}
|
||||
}
|
||||
|
||||
def fake_post(url, json, timeout):
|
||||
captured["url"] = url
|
||||
captured["json"] = json
|
||||
captured["timeout"] = timeout
|
||||
return _FakeResponse(
|
||||
{
|
||||
"model": "TimesFM",
|
||||
"model_id": "google/timesfm-2.5-200m-pytorch",
|
||||
"predictions": {
|
||||
"2026-03-15": 15.4,
|
||||
"2026-03-16": 15.9,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
monkeypatch.setenv("POLYWEATHER_TIMESFM_ENABLED", "true")
|
||||
monkeypatch.setenv("POLYWEATHER_TIMESFM_SERVICE_URL", "http://timesfm:8011")
|
||||
monkeypatch.setattr("src.analysis.deb_algorithm.load_history", fake_load_history)
|
||||
monkeypatch.setattr("src.models.timesfm_adapter.requests.post", fake_post)
|
||||
|
||||
result = predict_timesfm_daily(
|
||||
city_name="ankara",
|
||||
forecast_dates=["2026-03-15", "2026-03-16"],
|
||||
daily_model_forecasts={
|
||||
"2026-03-15": {"ECMWF": 15.0, "GFS": 16.0},
|
||||
"2026-03-16": {"ECMWF": 16.0, "GFS": 17.0},
|
||||
},
|
||||
)
|
||||
|
||||
assert captured["url"] == "http://timesfm:8011/predict/daily"
|
||||
assert captured["timeout"] == 12.0
|
||||
assert captured["json"]["series_frequency"] == "D"
|
||||
assert captured["json"]["series_kind"] == "actual_high"
|
||||
assert len(captured["json"]["series"]) == 14
|
||||
assert result["enabled"] is True
|
||||
assert result["reason"] == "ok"
|
||||
assert result["history_count"] == 14
|
||||
assert result["predictions"] == {
|
||||
"2026-03-15": 15.4,
|
||||
"2026-03-16": 15.9,
|
||||
}
|
||||
|
||||
|
||||
def test_predict_timesfm_daily_skips_when_history_is_insufficient(monkeypatch):
|
||||
def fake_load_history(_filepath):
|
||||
return {
|
||||
"ankara": {
|
||||
"2026-03-10": {"actual_high": 15.0},
|
||||
"2026-03-11": {"actual_high": 16.0},
|
||||
}
|
||||
}
|
||||
|
||||
monkeypatch.setenv("POLYWEATHER_TIMESFM_ENABLED", "true")
|
||||
monkeypatch.setenv("POLYWEATHER_TIMESFM_SERVICE_URL", "http://timesfm:8011")
|
||||
monkeypatch.setattr("src.analysis.deb_algorithm.load_history", fake_load_history)
|
||||
|
||||
result = predict_timesfm_daily(
|
||||
city_name="ankara",
|
||||
forecast_dates=["2026-03-15", "2026-03-16"],
|
||||
)
|
||||
|
||||
assert result["predictions"] == {}
|
||||
assert result["enabled"] is True
|
||||
assert result["reason"] == "insufficient_history"
|
||||
assert result["history_count"] == 2
|
||||
@@ -0,0 +1,25 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /srv/timesfm_service
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PIP_NO_CACHE_DIR=1 \
|
||||
TIMESFM_HOST=0.0.0.0 \
|
||||
TIMESFM_PORT=8011
|
||||
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends git \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
COPY timesfm_service/requirements.txt /tmp/requirements.txt
|
||||
RUN pip install --upgrade pip \
|
||||
&& pip install -r /tmp/requirements.txt
|
||||
|
||||
# Official upstream install path from google-research/timesfm README.
|
||||
RUN git clone --depth 1 https://github.com/google-research/timesfm.git /tmp/timesfm \
|
||||
&& pip install /tmp/timesfm
|
||||
|
||||
COPY timesfm_service /srv/timesfm_service
|
||||
|
||||
CMD ["sh", "-c", "uvicorn app:app --host 0.0.0.0 --port ${TIMESFM_PORT:-8011}"]
|
||||
@@ -0,0 +1,32 @@
|
||||
# PolyWeather TimesFM Service
|
||||
|
||||
This service isolates official `timesfm` inference from the main PolyWeather backend.
|
||||
|
||||
## What it does
|
||||
|
||||
- runs on Python 3.11+
|
||||
- installs official upstream `google-research/timesfm`
|
||||
- accepts recent `actual_high` or other temperature series payloads
|
||||
- returns point forecasts for the next 1 to 3 days
|
||||
|
||||
## Local run
|
||||
|
||||
```bash
|
||||
cd timesfm_service
|
||||
python -m venv .venv
|
||||
. .venv/bin/activate
|
||||
pip install -r requirements.txt
|
||||
git clone https://github.com/google-research/timesfm.git ../.timesfm-src
|
||||
pip install ../.timesfm-src
|
||||
uvicorn app:app --host 0.0.0.0 --port 8011
|
||||
```
|
||||
|
||||
## Docker
|
||||
|
||||
Build and run through the repo `docker-compose.yml`, or directly:
|
||||
|
||||
```bash
|
||||
docker compose --profile timesfm up --build polyweather_timesfm
|
||||
docker build -f timesfm_service/Dockerfile -t polyweather-timesfm .
|
||||
docker run --rm -p 8011:8011 polyweather-timesfm
|
||||
```
|
||||
@@ -0,0 +1,186 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
from typing import Dict, List
|
||||
|
||||
import numpy as np
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
DEFAULT_MODEL_ID = "google/timesfm-2.5-200m-pytorch"
|
||||
|
||||
|
||||
class SeriesPoint(BaseModel):
|
||||
timestamp: str
|
||||
value: float
|
||||
|
||||
|
||||
class DailyPredictRequest(BaseModel):
|
||||
city: str
|
||||
series_frequency: str = Field(default="D")
|
||||
series_kind: str = Field(default="actual_high")
|
||||
series: List[SeriesPoint]
|
||||
future_dates: List[str]
|
||||
daily_model_forecasts: Dict[str, Dict[str, float]] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class TimesFMPredictor:
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._model = None
|
||||
self._model_id = (
|
||||
str(os.getenv("TIMESFM_MODEL_ID", DEFAULT_MODEL_ID)).strip()
|
||||
or DEFAULT_MODEL_ID
|
||||
)
|
||||
self._max_context = int(os.getenv("TIMESFM_MAX_CONTEXT", "1024"))
|
||||
self._max_horizon = int(os.getenv("TIMESFM_MAX_HORIZON", "7"))
|
||||
self._normalize_inputs = str(
|
||||
os.getenv("TIMESFM_NORMALIZE_INPUTS", "true")
|
||||
).strip().lower() in {"1", "true", "yes", "on"}
|
||||
self._use_quantile_head = str(
|
||||
os.getenv("TIMESFM_USE_QUANTILE_HEAD", "true")
|
||||
).strip().lower() in {"1", "true", "yes", "on"}
|
||||
self._force_flip_invariance = str(
|
||||
os.getenv("TIMESFM_FORCE_FLIP_INVARIANCE", "true")
|
||||
).strip().lower() in {"1", "true", "yes", "on"}
|
||||
self._infer_is_positive = str(
|
||||
os.getenv("TIMESFM_INFER_IS_POSITIVE", "false")
|
||||
).strip().lower() in {"1", "true", "yes", "on"}
|
||||
self._fix_quantile_crossing = str(
|
||||
os.getenv("TIMESFM_FIX_QUANTILE_CROSSING", "true")
|
||||
).strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
@property
|
||||
def model_id(self) -> str:
|
||||
return self._model_id
|
||||
|
||||
def is_loaded(self) -> bool:
|
||||
return self._model is not None
|
||||
|
||||
def _ensure_loaded(self):
|
||||
if self._model is not None:
|
||||
return self._model
|
||||
|
||||
with self._lock:
|
||||
if self._model is not None:
|
||||
return self._model
|
||||
|
||||
import torch
|
||||
import timesfm
|
||||
|
||||
model_type = getattr(timesfm, "TimesFM_2p5_200M_torch", None)
|
||||
forecast_config_type = getattr(timesfm, "ForecastConfig", None)
|
||||
if model_type is None or forecast_config_type is None:
|
||||
raise RuntimeError("Unsupported official timesfm package layout.")
|
||||
|
||||
torch.set_float32_matmul_precision("high")
|
||||
model = model_type.from_pretrained(self._model_id)
|
||||
model.compile(
|
||||
forecast_config_type(
|
||||
max_context=self._max_context,
|
||||
max_horizon=self._max_horizon,
|
||||
normalize_inputs=self._normalize_inputs,
|
||||
use_continuous_quantile_head=self._use_quantile_head,
|
||||
force_flip_invariance=self._force_flip_invariance,
|
||||
infer_is_positive=self._infer_is_positive,
|
||||
fix_quantile_crossing=self._fix_quantile_crossing,
|
||||
)
|
||||
)
|
||||
self._model = model
|
||||
return self._model
|
||||
|
||||
def predict_daily(self, series: List[float], future_dates: List[str]) -> Dict[str, object]:
|
||||
model = self._ensure_loaded()
|
||||
horizon = len(future_dates)
|
||||
if horizon <= 0:
|
||||
raise ValueError("future_dates must not be empty")
|
||||
if horizon > self._max_horizon:
|
||||
raise ValueError(
|
||||
"future_dates exceeds configured max horizon: "
|
||||
f"{horizon} > {self._max_horizon}"
|
||||
)
|
||||
if len(series) < 8:
|
||||
raise ValueError("series must contain at least 8 points")
|
||||
|
||||
point_forecast, quantile_forecast = model.forecast(
|
||||
horizon=horizon,
|
||||
inputs=[np.asarray(series, dtype=np.float32)],
|
||||
)
|
||||
|
||||
point_row = point_forecast[0]
|
||||
predictions: Dict[str, float] = {}
|
||||
for index, date_str in enumerate(future_dates):
|
||||
if index >= len(point_row):
|
||||
break
|
||||
predictions[date_str] = round(float(point_row[index]), 1)
|
||||
|
||||
quantiles: Dict[str, Dict[str, float]] = {}
|
||||
try:
|
||||
if quantile_forecast is not None:
|
||||
q_arr = np.asarray(quantile_forecast)
|
||||
if (
|
||||
q_arr.ndim == 3
|
||||
and q_arr.shape[0] > 0
|
||||
and q_arr.shape[1] >= horizon
|
||||
and q_arr.shape[2] >= 10
|
||||
):
|
||||
for index, date_str in enumerate(future_dates):
|
||||
quantiles[date_str] = {
|
||||
"mean": round(float(q_arr[0, index, 0]), 1),
|
||||
"p10": round(float(q_arr[0, index, 1]), 1),
|
||||
"p50": round(float(q_arr[0, index, 5]), 1),
|
||||
"p90": round(float(q_arr[0, index, 9]), 1),
|
||||
}
|
||||
except Exception:
|
||||
quantiles = {}
|
||||
|
||||
return {
|
||||
"predictions": predictions,
|
||||
"quantiles": quantiles,
|
||||
}
|
||||
|
||||
|
||||
predictor = TimesFMPredictor()
|
||||
app = FastAPI(title="PolyWeather TimesFM Service", version="0.1.0")
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health() -> Dict[str, object]:
|
||||
return {
|
||||
"ok": True,
|
||||
"model_loaded": predictor.is_loaded(),
|
||||
"model_id": predictor.model_id,
|
||||
}
|
||||
|
||||
|
||||
@app.post("/predict/daily")
|
||||
def predict_daily(payload: DailyPredictRequest) -> Dict[str, object]:
|
||||
values = [float(point.value) for point in payload.series]
|
||||
future_dates = [
|
||||
str(date_str or "").strip()
|
||||
for date_str in payload.future_dates
|
||||
if str(date_str or "").strip()
|
||||
]
|
||||
if not future_dates:
|
||||
raise HTTPException(status_code=400, detail="future_dates must not be empty")
|
||||
if not values:
|
||||
raise HTTPException(status_code=400, detail="series must not be empty")
|
||||
|
||||
try:
|
||||
result = predictor.predict_daily(values, future_dates)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc))
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=500, detail=f"TimesFM inference failed: {exc}")
|
||||
|
||||
return {
|
||||
"model": "TimesFM",
|
||||
"model_id": predictor.model_id,
|
||||
"city": payload.city,
|
||||
"series_frequency": payload.series_frequency,
|
||||
"series_kind": payload.series_kind,
|
||||
"input_points": len(values),
|
||||
**result,
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
fastapi>=0.115,<1.0
|
||||
uvicorn[standard]>=0.30,<1.0
|
||||
numpy>=1.26,<3.0
|
||||
pydantic>=2.8,<3.0
|
||||
torch>=2.6,<3.0
|
||||
Reference in New Issue
Block a user