Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| """ | |
| 虫巢-200M HF Spaces入口 v5 | |
| 启动时下载训练权重 + 后台训练 + 前台API服务 | |
| """ | |
| import os | |
| import sys | |
| import json | |
| import threading | |
| import time | |
| BASE = os.path.dirname(os.path.abspath(__file__)) | |
| SRC = os.path.join(BASE, 'src') | |
| for p in [BASE, SRC, os.path.join(SRC, 'core'), os.path.join(SRC, 'training')]: | |
| if p not in sys.path: | |
| sys.path.insert(0, p) | |
| def download_weights(): | |
| """从HF Dataset下载模型权重""" | |
| repo = os.environ.get('HF_MODEL_REPO', '') | |
| token = os.environ.get('HF_TOKEN', '') | |
| if not repo: | |
| print("[HF] 未设置HF_MODEL_REPO，跳过权重下载") | |
| return False | |
| models_dir = os.path.join(BASE, 'models') | |
| os.makedirs(models_dir, exist_ok=True) | |
| try: | |
| from huggingface_hub import hf_hub_download | |
| weight_files = [ | |
| 'trained_200m_v2.npz', | |
| 'decoder_weights.npz', | |
| 'vocab75_clean_v2.json', | |
| 'vocab75_index_words.json', | |
| 'hippo_index.pkl', | |
| ] | |
| for f in weight_files: | |
| target = os.path.join(models_dir, f) | |
| if os.path.exists(target) and os.path.getsize(target) > 1000: | |
| print(f"[HF] {f} 已存在，跳过") | |
| continue | |
| print(f"[HF] 下载 {f}...") | |
| try: | |
| hf_hub_download( | |
| repo_id=repo, | |
| filename=f'models/{f}', | |
| token=token if token else None, | |
| local_dir=BASE, | |
| ) | |
| print(f"[HF] done {f}") | |
| except Exception as e: | |
| print(f"[HF] {f} 下载失败: {e}") | |
| print("[HF] 权重下载完成") | |
| return True | |
| except Exception as e: | |
| print(f"[HF] 权重下载失败: {e}") | |
| return False | |
| def run_training(): | |
| """后台训练线程""" | |
| time.sleep(30) # 等待权重下载 | |
| try: | |
| import subprocess | |
| proc = subprocess.Popen( | |
| [sys.executable, os.path.join(BASE, 'train_200m.py')], | |
| cwd=BASE, | |
| stdout=open(os.path.join(BASE, 'training.log'), 'w'), | |
| stderr=subprocess.STDOUT, | |
| ) | |
| proc.wait() | |
| except Exception as e: | |
| print(f'[训练] 异常: {e}') | |
| # 1. 下载权重 | |
| print("=" * 50) | |
| print(" 虫巢-200M HF Space 启动中...") | |
| print("=" * 50) | |
| download_weights() | |
| # 2. 启动训练线程 | |
| if os.environ.get('ENABLE_TRAINING', '0') == '1': | |
| train_thread = threading.Thread(target=run_training, daemon=True) | |
| train_thread.start() | |
| print('[训练] 后台训练已启动') | |
| # 3. FastAPI服务 | |
| from fastapi import FastAPI | |
| from fastapi.responses import HTMLResponse | |
| app = FastAPI(title="虫巢-200M") | |
| async def index(): | |
| status = {} | |
| try: | |
| with open(os.path.join(BASE, 'training_status.json')) as f: | |
| status = json.load(f) | |
| except: | |
| pass | |
| log_lines = [] | |
| try: | |
| with open(os.path.join(BASE, 'training.log')) as f: | |
| lines = f.readlines() | |
| log_lines = lines[-50:] | |
| except: | |
| log_lines = ['等待训练启动...'] | |
| return f"""<html><head><title>虫巢-200M 训练</title> | |
| <meta http-equiv="refresh" content="30"></head> | |
| <body style="font-family:monospace;max-width:900px;margin:40px auto;padding:0 20px;background:#1a1a2e;color:#e0e0e0"> | |
| <h1>🧠 虫巢-200M 训练监控</h1> | |
| <h2>状态</h2> | |
| <pre>{json.dumps(status, indent=2, ensure_ascii=False) if status else '等待中...'}</pre> | |
| <h2>训练日志 (最近50行, 30秒自动刷新)</h2> | |
| <pre style="background:#16213e;padding:15px;border-radius:8px;overflow-x:auto;font-size:13px">{''.join(log_lines[-50:])}</pre> | |
| </body></html>""" | |
| async def health(): | |
| return {"status": "ok", "model": "200M"} | |
| async def status(): | |
| try: | |
| with open(os.path.join(BASE, 'training_status.json')) as f: | |
| return json.load(f) | |
| except: | |
| return {"status": "starting"} | |
| if __name__ == '__main__': | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=7860) | |