#!/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")
@app.get("/", response_class=HTMLResponse)
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"""虫巢-200M 训练
🧠 虫巢-200M 训练监控
状态
{json.dumps(status, indent=2, ensure_ascii=False) if status else '等待中...'}
训练日志 (最近50行, 30秒自动刷新)
{''.join(log_lines[-50:])}
"""
@app.get("/health")
async def health():
return {"status": "ok", "model": "200M"}
@app.get("/status")
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)