Spaces:

Kurtobe
/

kronos-filter

Running

az30667176-sudo

fix: revert ZeroGPU (requires HF Pro); make @spaces.GPU no-op on free CPU

29841a6 7 months ago

8.53 kB

"""
Kronos Filter — HuggingFace Space (ZeroGPU)
Gradio app that wraps Kronos inference:
1. Visual UI (this page)
2. REST API at /call/predict (consumed by kronos-filter.vercel.app)
Runs on HF Spaces free CPU tier. @spaces.GPU decorator is kept as a no-op
so the code works on both CPU Spaces and (if upgraded to ZeroGPU via Pro)
GPU Spaces.
"""
from __future__ import annotations

import logging
import sys
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path

# Kronos model code lives at ./model/ in this Space
SPACE_ROOT = Path(__file__).resolve().parent
if str(SPACE_ROOT) not in sys.path:
    sys.path.insert(0, str(SPACE_ROOT))

import gradio as gr
import numpy as np
import torch

# @spaces.GPU decorator only useful on ZeroGPU (requires HF Pro).
# On free CPU tier we import it conditionally as a no-op.
try:
    import spaces
    SPACES_AVAILABLE = True
except ImportError:
    SPACES_AVAILABLE = False
    class _NoopGPU:
        def __call__(self, *a, **kw):
            def decorator(fn):
                return fn
            if len(a) == 1 and callable(a[0]):
                return a[0]
            return decorator
    class _NoopSpaces:
        GPU = _NoopGPU()
    spaces = _NoopSpaces()

from fetcher import fetch_aligned, build_future_timestamps
from predictor import KronosPathsPredictor
from probability_analyzer import analyze_ticker

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

# Predictor is lazily initialized on first GPU call, then cached
PREDICTOR: KronosPathsPredictor | None = None


def _ensure_predictor_on_device(device: str) -> KronosPathsPredictor:
    """
    Lazy-init predictor and ensure it's on the target device.
    With ZeroGPU, GPU is only accessible inside @spaces.GPU blocks, so we
    initialize on CPU first then migrate to GPU inside the decorator.
    """
    global PREDICTOR
    if PREDICTOR is None:
        logger.info(f"Initializing Kronos-small on {device}...")
        PREDICTOR = KronosPathsPredictor(model_size="small", device=device)
    elif PREDICTOR.device != device:
        logger.info(f"Moving Kronos from {PREDICTOR.device} to {device}...")
        PREDICTOR.model = PREDICTOR.model.to(device)
        PREDICTOR.tokenizer = PREDICTOR.tokenizer.to(device)
        PREDICTOR.device = device
    return PREDICTOR


@spaces.GPU(duration=180)
def run_prediction(
    tickers_raw: str,
    lookback: int,
    pred_len: int,
    samples: int,
    seed_raw: str,
) -> dict:
    """
    Main prediction pipeline — runs on ZeroGPU (A10G when available).
    Duration budget: 180s max per call.
    """
    # Parse tickers (comma or newline separated)
    tickers = []
    for chunk in tickers_raw.replace("\n", ",").split(","):
        t = chunk.strip().upper()
        if t and t not in tickers:
            tickers.append(t)
    if not tickers:
        raise gr.Error("Please provide at least one ticker.")
    if len(tickers) > 30:
        raise gr.Error("Max 30 tickers per request.")

    lookback = int(lookback)
    pred_len = int(pred_len)
    samples = int(samples)
    seed = None
    if seed_raw and seed_raw.strip().isdigit():
        seed = int(seed_raw.strip())

    logger.info(f"Request: tickers={tickers}, lookback={lookback}, pred_len={pred_len}, samples={samples}, seed={seed}")

    # 1. Fetch OHLCV (on CPU, no GPU needed)
    ok_list, failed = fetch_aligned(tickers, lookback=lookback, interval="1d")
    if not ok_list:
        raise gr.Error(f"Could not fetch any valid ticker data. Failed: {failed}")

    # 2. Future timestamps
    future_ts_list = [
        build_future_timestamps(td.timestamps.iloc[-1], pred_len, "1d")
        for td in ok_list
    ]

    # 3. Ensure model is on GPU (ZeroGPU allocates cuda:0 inside this decorator)
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    logger.info(f"Inference device: {device} (cuda available: {torch.cuda.is_available()})")
    predictor = _ensure_predictor_on_device(device)

    paths_list = predictor.predict_paths(
        df_list=[td.df for td in ok_list],
        x_timestamp_list=[td.timestamps for td in ok_list],
        y_timestamp_list=future_ts_list,
        pred_len=pred_len,
        T=1.0, top_p=0.9,
        sample_count=samples,
        verbose=False,
        seed=seed,
    )

    # 4. Compute metrics (CPU, fast)
    metrics_list = []
    for td, paths in zip(ok_list, paths_list):
        try:
            m = analyze_ticker(td.ticker, paths, initial_close=float(td.df["close"].iloc[-1]))
            metrics_list.append(m)
        except Exception as exc:
            logger.warning(f"[{td.ticker}] analysis failed: {exc}")

    # 5. Build JSON report (same shape as local CLI's latest.json)
    report = {
        "generated_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "config": {
            "lookback": lookback,
            "pred_len": pred_len,
            "samples": samples,
            "model_size": "small",
            "interval": "1d",
            "T": 1.0,
            "top_p": 0.9,
            "seed": seed,
            "device": device,
        },
        "failed_tickers": failed,
        "data_last_date": ok_list[0].timestamps.iloc[-1].isoformat() if ok_list else None,
        "tickers": sorted(
            [asdict(m) for m in metrics_list],
            key=lambda x: x["prob_up"],
            reverse=True,
        ),
    }

    # Path summaries for fan chart
    path_summaries = []
    for td, paths, ts_future in zip(ok_list, paths_list, future_ts_list):
        close_paths = paths[:, :, 3]
        path_summaries.append({
            "ticker": td.ticker,
            "history_dates": [d.isoformat() for d in td.timestamps.iloc[-60:]],
            "history_close": td.df["close"].iloc[-60:].tolist(),
            "future_dates": [d.isoformat() for d in ts_future],
            "median_path": np.median(close_paths, axis=0).tolist(),
            "p5_path": np.percentile(close_paths, 5, axis=0).tolist(),
            "p25_path": np.percentile(close_paths, 25, axis=0).tolist(),
            "p75_path": np.percentile(close_paths, 75, axis=0).tolist(),
            "p95_path": np.percentile(close_paths, 95, axis=0).tolist(),
        })
    report["path_summaries"] = path_summaries

    return report


# ----------------------------------------------------------------- UI

with gr.Blocks(title="Kronos Filter API") as demo:
    gr.Markdown(
        """
        # 📈 Kronos Filter — Probability API (ZeroGPU)
        Compute backend for [kronos-filter.vercel.app](https://kronos-filter.vercel.app).
        Fill in tickers to get probabilistic forecasts, or call programmatically via `/call/predict`.
        > ⚙️ **Hardware**: ZeroGPU (A10G, dynamically allocated).
        > ⚠️ Not financial advice. Results are **model-subjective probabilities**.
        """
    )
    with gr.Row():
        with gr.Column(scale=2):
            tickers_in = gr.Textbox(
                label="Tickers (comma or newline separated)",
                placeholder="MSFT, NVDA, AAPL",
                lines=3,
                value="MSFT, NVDA, AAPL",
            )
            with gr.Row():
                lookback_in = gr.Number(label="Lookback", value=400, minimum=50, maximum=512, precision=0)
                pred_len_in = gr.Number(label="Pred len", value=30, minimum=1, maximum=120, precision=0)
                samples_in = gr.Number(label="Samples", value=30, minimum=5, maximum=60, precision=0)
                seed_in = gr.Textbox(label="Seed (blank = random)", value="42")
            run_btn = gr.Button("🚀 Predict", variant="primary")
        with gr.Column(scale=3):
            output_json = gr.JSON(label="Probability report", open=False)

    run_btn.click(
        fn=run_prediction,
        inputs=[tickers_in, lookback_in, pred_len_in, samples_in, seed_in],
        outputs=output_json,
        api_name="predict",
    )

    gr.Markdown(
        """
        ---
        **Model**: Kronos-small (24.7M params, 512 context) · ZeroGPU A10G ·
        First request cold-start ~15s (model load onto GPU), subsequent calls ~5-15s.
        Daily GPU quota for free HF accounts: ~5-10 min.
        Source: [github.com/az30667176-sudo/kronos-filter](https://github.com/az30667176-sudo/kronos-filter)
        """
    )


demo.queue(default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)

	"""
	Kronos Filter — HuggingFace Space (ZeroGPU)

	Gradio app that wraps Kronos inference:
	1. Visual UI (this page)
	2. REST API at /call/predict (consumed by kronos-filter.vercel.app)

	Runs on HF Spaces free CPU tier. @spaces.GPU decorator is kept as a no-op
	so the code works on both CPU Spaces and (if upgraded to ZeroGPU via Pro)
	GPU Spaces.
	"""
	from __future__ import annotations

	import logging
	import sys
	from dataclasses import asdict
	from datetime import datetime, timezone
	from pathlib import Path

	# Kronos model code lives at ./model/ in this Space
	SPACE_ROOT = Path(__file__).resolve().parent
	if str(SPACE_ROOT) not in sys.path:
	sys.path.insert(0, str(SPACE_ROOT))

	import gradio as gr
	import numpy as np
	import torch

	# @spaces.GPU decorator only useful on ZeroGPU (requires HF Pro).
	# On free CPU tier we import it conditionally as a no-op.
	try:
	import spaces
	SPACES_AVAILABLE = True
	except ImportError:
	SPACES_AVAILABLE = False
	class _NoopGPU:
	def __call__(self, a, *kw):
	def decorator(fn):
	return fn
	if len(a) == 1 and callable(a[0]):
	return a[0]
	return decorator
	class _NoopSpaces:
	GPU = _NoopGPU()
	spaces = _NoopSpaces()

	from fetcher import fetch_aligned, build_future_timestamps
	from predictor import KronosPathsPredictor
	from probability_analyzer import analyze_ticker

	logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
	logger = logging.getLogger(__name__)

	# Predictor is lazily initialized on first GPU call, then cached
	PREDICTOR: KronosPathsPredictor \| None = None


	def _ensure_predictor_on_device(device: str) -> KronosPathsPredictor:
	"""
	Lazy-init predictor and ensure it's on the target device.
	With ZeroGPU, GPU is only accessible inside @spaces.GPU blocks, so we
	initialize on CPU first then migrate to GPU inside the decorator.
	"""
	global PREDICTOR
	if PREDICTOR is None:
	logger.info(f"Initializing Kronos-small on {device}...")
	PREDICTOR = KronosPathsPredictor(model_size="small", device=device)
	elif PREDICTOR.device != device:
	logger.info(f"Moving Kronos from {PREDICTOR.device} to {device}...")
	PREDICTOR.model = PREDICTOR.model.to(device)
	PREDICTOR.tokenizer = PREDICTOR.tokenizer.to(device)
	PREDICTOR.device = device
	return PREDICTOR


	@spaces.GPU(duration=180)
	def run_prediction(
	tickers_raw: str,
	lookback: int,
	pred_len: int,
	samples: int,
	seed_raw: str,
	) -> dict:
	"""
	Main prediction pipeline — runs on ZeroGPU (A10G when available).
	Duration budget: 180s max per call.
	"""
	# Parse tickers (comma or newline separated)
	tickers = []
	for chunk in tickers_raw.replace("\n", ",").split(","):
	t = chunk.strip().upper()
	if t and t not in tickers:
	tickers.append(t)
	if not tickers:
	raise gr.Error("Please provide at least one ticker.")
	if len(tickers) > 30:
	raise gr.Error("Max 30 tickers per request.")

	lookback = int(lookback)
	pred_len = int(pred_len)
	samples = int(samples)
	seed = None
	if seed_raw and seed_raw.strip().isdigit():
	seed = int(seed_raw.strip())

	logger.info(f"Request: tickers={tickers}, lookback={lookback}, pred_len={pred_len}, samples={samples}, seed={seed}")

	# 1. Fetch OHLCV (on CPU, no GPU needed)
	ok_list, failed = fetch_aligned(tickers, lookback=lookback, interval="1d")
	if not ok_list:
	raise gr.Error(f"Could not fetch any valid ticker data. Failed: {failed}")

	# 2. Future timestamps
	future_ts_list = [
	build_future_timestamps(td.timestamps.iloc[-1], pred_len, "1d")
	for td in ok_list
	]

	# 3. Ensure model is on GPU (ZeroGPU allocates cuda:0 inside this decorator)
	device = "cuda:0" if torch.cuda.is_available() else "cpu"
	logger.info(f"Inference device: {device} (cuda available: {torch.cuda.is_available()})")
	predictor = _ensure_predictor_on_device(device)

	paths_list = predictor.predict_paths(
	df_list=[td.df for td in ok_list],
	x_timestamp_list=[td.timestamps for td in ok_list],
	y_timestamp_list=future_ts_list,
	pred_len=pred_len,
	T=1.0, top_p=0.9,
	sample_count=samples,
	verbose=False,
	seed=seed,
	)

	# 4. Compute metrics (CPU, fast)
	metrics_list = []
	for td, paths in zip(ok_list, paths_list):
	try:
	m = analyze_ticker(td.ticker, paths, initial_close=float(td.df["close"].iloc[-1]))
	metrics_list.append(m)
	except Exception as exc:
	logger.warning(f"[{td.ticker}] analysis failed: {exc}")

	# 5. Build JSON report (same shape as local CLI's latest.json)
	report = {
	"generated_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
	"config": {
	"lookback": lookback,
	"pred_len": pred_len,
	"samples": samples,
	"model_size": "small",
	"interval": "1d",
	"T": 1.0,
	"top_p": 0.9,
	"seed": seed,
	"device": device,
	},
	"failed_tickers": failed,
	"data_last_date": ok_list[0].timestamps.iloc[-1].isoformat() if ok_list else None,
	"tickers": sorted(
	[asdict(m) for m in metrics_list],
	key=lambda x: x["prob_up"],
	reverse=True,
	),
	}

	# Path summaries for fan chart
	path_summaries = []
	for td, paths, ts_future in zip(ok_list, paths_list, future_ts_list):
	close_paths = paths[:, :, 3]
	path_summaries.append({
	"ticker": td.ticker,
	"history_dates": [d.isoformat() for d in td.timestamps.iloc[-60:]],
	"history_close": td.df["close"].iloc[-60:].tolist(),
	"future_dates": [d.isoformat() for d in ts_future],
	"median_path": np.median(close_paths, axis=0).tolist(),
	"p5_path": np.percentile(close_paths, 5, axis=0).tolist(),
	"p25_path": np.percentile(close_paths, 25, axis=0).tolist(),
	"p75_path": np.percentile(close_paths, 75, axis=0).tolist(),
	"p95_path": np.percentile(close_paths, 95, axis=0).tolist(),
	})
	report["path_summaries"] = path_summaries

	return report


	# ----------------------------------------------------------------- UI

	with gr.Blocks(title="Kronos Filter API") as demo:
	gr.Markdown(
	"""
	# 📈 Kronos Filter — Probability API (ZeroGPU)

	Compute backend for [kronos-filter.vercel.app](https://kronos-filter.vercel.app).
	Fill in tickers to get probabilistic forecasts, or call programmatically via `/call/predict`.

	> ⚙️ Hardware: ZeroGPU (A10G, dynamically allocated).
	> ⚠️ Not financial advice. Results are model-subjective probabilities.
	"""
	)
	with gr.Row():
	with gr.Column(scale=2):
	tickers_in = gr.Textbox(
	label="Tickers (comma or newline separated)",
	placeholder="MSFT, NVDA, AAPL",
	lines=3,
	value="MSFT, NVDA, AAPL",
	)
	with gr.Row():
	lookback_in = gr.Number(label="Lookback", value=400, minimum=50, maximum=512, precision=0)
	pred_len_in = gr.Number(label="Pred len", value=30, minimum=1, maximum=120, precision=0)
	samples_in = gr.Number(label="Samples", value=30, minimum=5, maximum=60, precision=0)
	seed_in = gr.Textbox(label="Seed (blank = random)", value="42")
	run_btn = gr.Button("🚀 Predict", variant="primary")
	with gr.Column(scale=3):
	output_json = gr.JSON(label="Probability report", open=False)

	run_btn.click(
	fn=run_prediction,
	inputs=[tickers_in, lookback_in, pred_len_in, samples_in, seed_in],
	outputs=output_json,
	api_name="predict",
	)

	gr.Markdown(
	"""
	---
	Model: Kronos-small (24.7M params, 512 context) · ZeroGPU A10G ·
	First request cold-start ~15s (model load onto GPU), subsequent calls ~5-15s.
	Daily GPU quota for free HF accounts: ~5-10 min.
	Source: [github.com/az30667176-sudo/kronos-filter](https://github.com/az30667176-sudo/kronos-filter)
	"""
	)


	demo.queue(default_concurrency_limit=1)

	if __name__ == "__main__":
	demo.launch(server_name="0.0.0.0", server_port=7860, share=True)