wehub-resource-sync
262 lines
8.8 KiB
Python
262 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Gemma4 教程冒烟测试（不默认加载整模权重，避免 OOM；可按需开启网络拉取 Processor）。
|
|
|
|
用法:
|
|
/path/to/self-llm/bin/python verify_gemma4_tutorials.py
|
|
GEMMA_PULL_PROCESSOR=1 ... # 从 Hub 拉取 Gemma4Processor（需已安装 torchvision，与 torch 同 CUDA 版本）
|
|
|
|
说明：本脚本不做整模加载/推理；真机部署请参考各 .md 并自行下载权重。
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
from collections import UserDict
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
DATASET = REPO_ROOT / "dataset" / "huanhuan.json"
|
|
MODEL_ID = "google/gemma-4-E4B-it"
|
|
|
|
|
|
def _ok(name: str, detail: str = "") -> None:
|
|
print(f"[PASS] {name}" + (f" — {detail}" if detail else ""))
|
|
|
|
|
|
def _fail(name: str, err: BaseException) -> None:
|
|
print(f"[FAIL] {name}: {err}")
|
|
traceback.print_exc()
|
|
|
|
|
|
def test_imports() -> bool:
|
|
import numpy
|
|
import torch
|
|
import transformers
|
|
from transformers import AutoModelForMultimodalLM, AutoProcessor
|
|
|
|
assert hasattr(torch, "cuda")
|
|
_ok(
|
|
"import 栈",
|
|
f"numpy={numpy.__version__}, torch={torch.__version__}, "
|
|
f"cuda={torch.cuda.is_available()}, transformers={transformers.__version__}",
|
|
)
|
|
_ = AutoModelForMultimodalLM
|
|
_ = AutoProcessor
|
|
return True
|
|
|
|
|
|
def test_pull_processor() -> bool:
|
|
if os.environ.get("GEMMA_PULL_PROCESSOR", "").strip() not in ("1", "true", "yes"):
|
|
print("[SKIP] Processor Hub 拉取（设置 GEMMA_PULL_PROCESSOR=1 可开启）")
|
|
return True
|
|
from transformers import AutoProcessor
|
|
|
|
t0 = time.time()
|
|
proc = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
|
|
dt = time.time() - t0
|
|
_ok("AutoProcessor.from_pretrained", f"{type(proc).__name__}, {dt:.1f}s")
|
|
return True
|
|
|
|
|
|
def test_fastapi_smoke_with_mock() -> bool:
|
|
"""教程 01 路由逻辑：用 Mock 模型验证请求/响应（避免真推理）。"""
|
|
from contextlib import asynccontextmanager
|
|
|
|
from fastapi import Body, FastAPI, HTTPException
|
|
from fastapi.testclient import TestClient
|
|
from pydantic import BaseModel, Field, model_validator
|
|
from typing import List, Literal, Optional
|
|
|
|
class ContentItem(BaseModel):
|
|
type: Literal["text", "image"]
|
|
text: Optional[str] = Field(None, description="文本")
|
|
image: Optional[str] = Field(None, description="图片 URL 或 base64")
|
|
|
|
@model_validator(mode="after")
|
|
def _v(self):
|
|
if self.type == "text" and not (self.text and str(self.text).strip()):
|
|
raise ValueError("文本类型必须提供 text 字段")
|
|
if self.type == "image":
|
|
img = self.image or ""
|
|
if not img.startswith(("http://", "https://", "data:image")):
|
|
raise ValueError("图片必须是有效的 URL 或 base64(data:image)")
|
|
return self
|
|
|
|
class Message(BaseModel):
|
|
role: Literal["system", "user", "assistant"]
|
|
content: List[ContentItem]
|
|
|
|
class ProcessRequest(BaseModel):
|
|
messages: List[Message] = Field(..., min_length=1)
|
|
max_new_tokens: int = Field(1000, ge=10, le=4096)
|
|
|
|
class ProcessResponse(BaseModel):
|
|
response: str
|
|
status: int
|
|
time: int
|
|
processing_time: float
|
|
tokens_generated: int
|
|
|
|
import torch
|
|
|
|
class _Batch(UserDict):
|
|
def to(self, _device):
|
|
return self
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.device = "cpu"
|
|
mock_model.generate = MagicMock(return_value=torch.tensor([[1, 2, 99, 100]]))
|
|
|
|
mock_processor = MagicMock()
|
|
mock_processor.apply_chat_template = MagicMock(
|
|
return_value=_Batch(
|
|
input_ids=torch.tensor([[1, 2]]),
|
|
attention_mask=torch.tensor([[1, 1]]),
|
|
)
|
|
)
|
|
mock_processor.decode = MagicMock(return_value="<mock>")
|
|
mock_processor.parse_response = MagicMock(side_effect=Exception("no parse"))
|
|
|
|
model_ref = {"m": mock_model, "p": mock_processor}
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
yield
|
|
|
|
app = FastAPI(lifespan=lifespan)
|
|
DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."
|
|
|
|
def _normalize_content_items(items: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
out = []
|
|
for it in items:
|
|
if it.get("type") == "text":
|
|
out.append({"type": "text", "text": it.get("text") or ""})
|
|
elif it.get("type") == "image":
|
|
img = it.get("image") or it.get("url")
|
|
if not img:
|
|
continue
|
|
if str(img).startswith(("http://", "https://")):
|
|
out.append({"type": "image", "url": img})
|
|
else:
|
|
out.append({"type": "image", "image": img})
|
|
return out
|
|
|
|
@app.post("/chat/completions", response_model=ProcessResponse)
|
|
async def generate_response(chat: ProcessRequest = Body(...)):
|
|
start_time = time.time()
|
|
try:
|
|
model, processor = model_ref["m"], model_ref["p"]
|
|
processed_messages = []
|
|
system_prompt = DEFAULT_SYSTEM_PROMPT
|
|
for msg in chat.messages:
|
|
if msg.role == "system":
|
|
system_prompt = " ".join([item.text or "" for item in msg.content if item.type == "text"])
|
|
else:
|
|
d = msg.model_dump()
|
|
d["content"] = _normalize_content_items(d["content"])
|
|
processed_messages.append(d)
|
|
messages = [
|
|
{"role": "system", "content": [{"type": "text", "text": system_prompt}]},
|
|
*processed_messages,
|
|
]
|
|
inputs = processor.apply_chat_template(
|
|
messages,
|
|
add_generation_prompt=True,
|
|
tokenize=True,
|
|
return_tensors="pt",
|
|
return_dict=True,
|
|
).to(model.device)
|
|
input_len = inputs["input_ids"].shape[-1]
|
|
max_token_num = min(4096, int(chat.max_new_tokens))
|
|
with torch.inference_mode():
|
|
generation = model.generate(**inputs, max_new_tokens=max_token_num, do_sample=False)
|
|
response_ids = generation[0][input_len:]
|
|
raw = processor.decode(response_ids, skip_special_tokens=False)
|
|
try:
|
|
parsed = processor.parse_response(raw)
|
|
decoded = parsed.get("content", raw) if isinstance(parsed, dict) else raw
|
|
except Exception:
|
|
decoded = processor.decode(response_ids, skip_special_tokens=True)
|
|
return ProcessResponse(
|
|
response=str(decoded),
|
|
status=200,
|
|
time=int(time.time()),
|
|
processing_time=time.time() - start_time,
|
|
tokens_generated=int(len(response_ids)),
|
|
)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
payload = {
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "你好，只做连通性测试"},
|
|
],
|
|
}
|
|
],
|
|
"max_new_tokens": 64,
|
|
}
|
|
with TestClient(app) as client:
|
|
r = client.post("/chat/completions", json=payload)
|
|
assert r.status_code == 200, r.text
|
|
body = r.json()
|
|
assert "response" in body
|
|
_ok("FastAPI /chat/completions (Mock)", f"tokens_generated={body.get('tokens_generated')}")
|
|
return True
|
|
|
|
|
|
def test_lora_dataset() -> bool:
|
|
if not DATASET.is_file():
|
|
print(f"[SKIP] 数据集不存在: {DATASET}")
|
|
return True
|
|
raw = json.loads(DATASET.read_text(encoding="utf-8"))
|
|
assert isinstance(raw, list) and len(raw) > 0
|
|
first = raw[0]
|
|
for k in ("instruction", "input", "output"):
|
|
assert k in first
|
|
_ok("05 LoRA 数据集", f"{DATASET.name} 条数={len(raw)}")
|
|
return True
|
|
|
|
|
|
def test_evalscope_import() -> bool:
|
|
from evalscope.config import TaskConfig
|
|
from evalscope.run import run_task
|
|
|
|
_ = TaskConfig
|
|
_ = run_task
|
|
_ok("evalscope 导入", "TaskConfig / run_task 可用（完整评测需 Ollama 等服务）")
|
|
return True
|
|
|
|
|
|
def main() -> int:
|
|
tests = [
|
|
("环境导入", test_imports),
|
|
("05 嬛嬛数据集", test_lora_dataset),
|
|
("01 FastAPI Mock", test_fastapi_smoke_with_mock),
|
|
("04 evalscope", test_evalscope_import),
|
|
("Hub Processor", test_pull_processor),
|
|
]
|
|
failed = 0
|
|
for name, fn in tests:
|
|
try:
|
|
fn()
|
|
except Exception as e:
|
|
failed += 1
|
|
_fail(name, e)
|
|
print("---")
|
|
print(f"完成: {len(tests) - failed}/{len(tests)} 通过")
|
|
return 1 if failed else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|