|
|
from __future__ import annotations |
|
|
import os |
|
|
import time |
|
|
import random |
|
|
from typing import Dict, Any, List, Literal, Optional |
|
|
import httpx |
|
|
|
|
|
HF_API_MODELS = "https://huggingface.co/api/models" |
|
|
HF_API_DATASETS = "https://huggingface.co/api/datasets" |
|
|
|
|
|
REFRESH_INTERVAL_SEC = int(os.getenv("HF_REGISTRY_REFRESH_SEC", "21600")) |
|
|
HTTP_TIMEOUT = float(os.getenv("HF_HTTP_TIMEOUT", "8.0")) |
|
|
|
|
|
_SEED_MODELS = [ |
|
|
"ElKulako/cryptobert", |
|
|
"kk08/CryptoBERT", |
|
|
] |
|
|
_SEED_DATASETS = [ |
|
|
"linxy/CryptoCoin", |
|
|
"WinkingFace/CryptoLM-Bitcoin-BTC-USDT", |
|
|
"WinkingFace/CryptoLM-Ethereum-ETH-USDT", |
|
|
"WinkingFace/CryptoLM-Solana-SOL-USDT", |
|
|
"WinkingFace/CryptoLM-Ripple-XRP-USDT", |
|
|
] |
|
|
|
|
|
|
|
|
class HFRegistry: |
|
|
def __init__(self) -> None: |
|
|
self.models: Dict[str, Dict[str, Any]] = {} |
|
|
self.datasets: Dict[str, Dict[str, Any]] = {} |
|
|
self.last_refresh: float = 0.0 |
|
|
self.fail_reason: Optional[str] = None |
|
|
|
|
|
async def _hf_json(self, url: str, params: Dict[str, Any]) -> Any: |
|
|
async with httpx.AsyncClient(timeout=HTTP_TIMEOUT) as client: |
|
|
r = await client.get(url, params=params) |
|
|
r.raise_for_status() |
|
|
return r.json() |
|
|
|
|
|
async def refresh(self) -> Dict[str, Any]: |
|
|
try: |
|
|
for name in _SEED_MODELS: |
|
|
self.models.setdefault(name, {"id": name, "source": "seed", "pipeline_tag": "sentiment-analysis"}) |
|
|
for name in _SEED_DATASETS: |
|
|
self.datasets.setdefault(name, {"id": name, "source": "seed"}) |
|
|
|
|
|
q_crypto = {"search": "crypto", "limit": 50} |
|
|
q_sent = {"pipeline_tag": "sentiment-analysis", "search": "crypto", "limit": 50} |
|
|
|
|
|
models = await self._hf_json(HF_API_MODELS, q_sent) |
|
|
for m in models or []: |
|
|
mid = m.get("modelId") or m.get("id") or m.get("name") |
|
|
if not mid: continue |
|
|
self.models[mid] = { |
|
|
"id": mid, |
|
|
"pipeline_tag": m.get("pipeline_tag"), |
|
|
"likes": m.get("likes"), |
|
|
"downloads": m.get("downloads"), |
|
|
"tags": m.get("tags") or [], |
|
|
"source": "hub", |
|
|
} |
|
|
|
|
|
datasets = await self._hf_json(HF_API_DATASETS, q_crypto) |
|
|
for d in datasets or []: |
|
|
did = d.get("id") or d.get("name") |
|
|
if not did: continue |
|
|
self.datasets[did] = { |
|
|
"id": did, |
|
|
"likes": d.get("likes"), |
|
|
"downloads": d.get("downloads"), |
|
|
"tags": d.get("tags") or [], |
|
|
"source": "hub", |
|
|
} |
|
|
|
|
|
self.last_refresh = time.time() |
|
|
self.fail_reason = None |
|
|
return {"ok": True, "models": len(self.models), "datasets": len(self.datasets)} |
|
|
except Exception as e: |
|
|
self.fail_reason = str(e) |
|
|
return {"ok": False, "error": self.fail_reason, "models": len(self.models), "datasets": len(self.datasets)} |
|
|
|
|
|
def list(self, kind: Literal["models","datasets"]="models") -> List[Dict[str, Any]]: |
|
|
return list(self.models.values()) if kind == "models" else list(self.datasets.values()) |
|
|
|
|
|
def health(self) -> Dict[str, Any]: |
|
|
age = time.time() - (self.last_refresh or 0) |
|
|
return { |
|
|
"ok": self.last_refresh > 0 and (self.fail_reason is None), |
|
|
"last_refresh_epoch": self.last_refresh, |
|
|
"age_sec": age, |
|
|
"fail_reason": self.fail_reason, |
|
|
"counts": {"models": len(self.models), "datasets": len(self.datasets)}, |
|
|
"interval_sec": REFRESH_INTERVAL_SEC, |
|
|
} |
|
|
|
|
|
|
|
|
REGISTRY = HFRegistry() |
|
|
|
|
|
|
|
|
async def periodic_refresh(loop_sleep: int = REFRESH_INTERVAL_SEC) -> None: |
|
|
await REGISTRY.refresh() |
|
|
await _sleep(int(loop_sleep * random.uniform(0.5, 0.9))) |
|
|
while True: |
|
|
await REGISTRY.refresh() |
|
|
await _sleep(loop_sleep) |
|
|
|
|
|
|
|
|
async def _sleep(sec: int) -> None: |
|
|
import asyncio |
|
|
try: |
|
|
await asyncio.sleep(sec) |
|
|
except Exception: |
|
|
pass |
|
|
|
