Spaces:
Paused
Paused
File size: 5,110 Bytes
d2aafc6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | import sys
import os
import time
import contextlib
import random
import numpy as np
import torch
import tiktoken
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from contextlib import asynccontextmanager
from train_glm5 import GLM5ForCausalLM, GLM5Config
# Ensure PyTorch unpickler can resolve GLM5Config if saved under __main__
sys.modules["__main__"].GLM5Config = GLM5Config
# Global model state
MODEL_STATE = {
"model": None,
"tokenizer": None,
"device": "cpu",
"checkpoint_path": None,
"loaded": False,
}
@asynccontextmanager
async def lifespan(app: FastAPI):
"""
FastAPI Lifespan Context Manager.
Loads the trained Nano-GLM checkpoint into CPU memory once when the server starts.
"""
device = "cpu"
out_dir = "out_glm5"
# Priority order for checkpoint loading
candidate_ckpts = ["model_inference.pt", "ckpt_best.pt", "ckpt.pt"]
selected_ckpt = None
for ckpt_name in candidate_ckpts:
path = os.path.join(out_dir, ckpt_name)
if os.path.exists(path):
selected_ckpt = path
break
if selected_ckpt is None:
print(" [WARNING] No model checkpoint found. Serving dummy state.")
yield
return
print(f" [STARTUP] Loading Nano-GLM model from {selected_ckpt}...")
checkpoint = torch.load(selected_ckpt, map_location=device, weights_only=False)
config = checkpoint["config"]
model = GLM5ForCausalLM(config)
model.load_state_dict(checkpoint["model"])
model.to(device)
model.eval()
tokenizer = tiktoken.get_encoding("gpt2")
MODEL_STATE["model"] = model
MODEL_STATE["tokenizer"] = tokenizer
MODEL_STATE["device"] = device
MODEL_STATE["checkpoint_path"] = selected_ckpt
MODEL_STATE["loaded"] = True
print(f" [STARTUP] Model successfully loaded on {device}!")
yield
print(" [SHUTDOWN] Cleaning up model state...")
MODEL_STATE.clear()
app = FastAPI(
title="Nano-GLM (GLM-5.2 Baby 120M) API",
description="FastAPI Backend for serving the lightweight from-scratch GLM-5.2 MoE model.",
version="1.0.0",
lifespan=lifespan,
)
# Enable CORS for frontend integration
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
class GenerateRequest(BaseModel):
prompt: str = Field(default="In conclusion,", description="Input prompt for text generation")
max_new_tokens: int = Field(default=100, ge=1, le=500, description="Maximum new tokens to generate")
temperature: float = Field(default=0.8, ge=0.0, le=2.0, description="Sampling temperature (0.0 = greedy)")
top_k: int = Field(default=50, ge=0, le=200, description="Top-k sampling threshold (0 = disable)")
seed: int = Field(default=42, description="Random seed for reproducibility (-1 for random)")
class GenerateResponse(BaseModel):
prompt: str
generated_text: str
num_tokens: int
inference_time_sec: float
@app.get("/info")
def read_root():
return {
"title": "Nano-GLM API",
"status": "online" if MODEL_STATE["loaded"] else "model_not_loaded",
"docs_url": "/docs",
}
@app.get("/health")
def health_check():
return {
"status": "healthy" if MODEL_STATE["loaded"] else "degraded",
"model_loaded": MODEL_STATE["loaded"],
"checkpoint": MODEL_STATE["checkpoint_path"],
"device": MODEL_STATE["device"],
}
@app.post("/generate", response_model=GenerateResponse)
def generate_text(req: GenerateRequest):
if not MODEL_STATE["loaded"]:
raise HTTPException(status_code=503, detail="Model is not loaded. Train the model first or export a checkpoint.")
model = MODEL_STATE["model"]
enc = MODEL_STATE["tokenizer"]
device = MODEL_STATE["device"]
# Enforce seed reproducibility if specified
if req.seed != -1:
torch.manual_seed(req.seed)
np.random.seed(req.seed)
random.seed(req.seed)
t0 = time.time()
# Tokenize input prompt
tokens = enc.encode(req.prompt)
if len(tokens) == 0:
tokens = [enc.eot_token]
idx = torch.tensor([tokens], dtype=torch.long, device=device)
# Generation context
with torch.no_grad():
output = model.generate(
idx,
max_new_tokens=req.max_new_tokens,
temperature=req.temperature,
top_k=req.top_k if req.top_k > 0 else None,
)
t1 = time.time()
inference_time = t1 - t0
generated_tokens = output[0].tolist()
generated_text = enc.decode(generated_tokens)
return GenerateResponse(
prompt=req.prompt,
generated_text=generated_text,
num_tokens=len(generated_tokens),
inference_time_sec=round(inference_time, 4),
)
# Mount static web frontend at root (if static directory exists)
if os.path.exists("static"):
app.mount("/", StaticFiles(directory="static", html=True), name="static")
|