| |
| """s4_openai_api.py — samai-4b OpenAI 兼容 API (Flask, port 7862, key=1234). |
| |
| Endpoints: |
| GET /v1/models |
| POST /v1/chat/completions {model, messages, temperature, top_p, max_tokens, |
| stream, enable_thinking} (think -> reasoning_content) |
| GET /health |
| Auth: Authorization: Bearer 1234 (或 x-api-key: 1234) |
| |
| 示例: |
| curl http://127.0.0.1:7862/v1/chat/completions \\ |
| -H "Authorization: Bearer 1234" -H "Content-Type: application/json" \\ |
| -d '{"model":"samai-4b","messages":[{"role":"user","content":"你好"}]}' |
| |
| 协议: 默认强制思考 (enable_thinking=true, 与 SFT 训练格式一致); |
| 思考文本放 message.reasoning_content (DeepSeek-R1 风格), 正文放 message.content。 |
| 上下文: config.max_position_embeddings (200k) 为硬上限, 超限返回 400。 |
| 注意 T4 16GB 显存下 KV 实际可服务约 4~5 万 token 长输入。 |
| """ |
| import json |
| import os |
| import time |
| import threading |
| import uuid |
|
|
| import torch |
| from flask import Flask, request, jsonify, Response |
| from transformers import (AutoModelForCausalLM, AutoTokenizer, |
| StoppingCriteria, StoppingCriteriaList) |
|
|
| MODEL_DIR = os.environ.get("S4_MODEL_DIR", "/content/samai-4b-sft") |
| PORT = int(os.environ.get("S4_API_PORT", "7862")) |
| API_KEY = os.environ.get("S4_API_KEY", "1234") |
| MODEL_NAME = "samai-4b" |
| DEFAULT_MAX_TOKENS = 1024 |
| HARD_MAX_TOKENS = 8192 |
| EOS_IDS = [1] |
|
|
| STATE = {"loaded": False, "error": None, "t0": time.time()} |
| LOCK = threading.Lock() |
| MODEL = {"tok": None, "m": None, "ctx": 200000} |
|
|
| app = Flask(__name__) |
|
|
|
|
| class AntiLoop(StoppingCriteria): |
| """末尾片段(L=3..16 token)连续重复 >=3 次 -> 复读退化, 提前截断.""" |
| def __call__(self, input_ids, scores, **kwargs): |
| ids = input_ids[0].tolist() |
| tail = ids[-64:] |
| if len(tail) < 9: |
| return False |
| for L in range(3, 17): |
| if len(tail) < 3 * L: |
| break |
| seg = tail[-L:] |
| if seg == tail[-2 * L:-L] == tail[-3 * L:-2 * L]: |
| return True |
| return False |
|
|
|
|
| def build(): |
| global MODEL, STATE |
| print("[build] loading tokenizer ...", flush=True) |
| tok = AutoTokenizer.from_pretrained(MODEL_DIR) |
| print("[build] loading model fp16 -> cuda ...", flush=True) |
| m = AutoModelForCausalLM.from_pretrained( |
| MODEL_DIR, trust_remote_code=True, dtype=torch.float16).cuda().eval() |
| ctx = int(getattr(m.config, "max_position_embeddings", 200000)) |
| MODEL.update(tok=tok, m=m, ctx=ctx) |
| STATE["loaded"] = True |
| print(f"[build] DONE ctx={ctx}", flush=True) |
|
|
|
|
| MARKS = ["<think>", "</think>", "<|User|>", "<|Bot|>", "<|System|>", "<|Tool|>", |
| "<｜start▁of▁sentence｜>", "<｜end▁of▁sentence｜>", "<｜▁pad▁｜>", |
| "<｜start▁of▁text｜>", "<｜end▁of▁text｜>", "<unk>"] |
|
|
|
|
| def clean(s): |
| for mk in MARKS: |
| s = s.replace(mk, "") |
| return s.strip() |
|
|
|
|
| def split_think(text): |
| text = text.split("<｜end▁of▁sentence｜>")[0] |
| if "<think>" in text and "</think>" in text: |
| _, b = text.split("<think>", 1) |
| th, rest = b.split("</think>", 1) |
| return th.strip(), clean(rest) |
| if "<think>" in text: |
| return text.split("<think>", 1)[1].strip(), "" |
| if "</think>" in text: |
| th, rest = text.split("</think>", 1) |
| return th.strip(), clean(rest) |
| return "", clean(text) |
|
|
|
|
| def build_inputs(msgs, enable_think): |
| tok = MODEL["tok"] |
| text = tok.apply_chat_template(msgs, add_generation_prompt=True, |
| tokenize=False, enable_thinking=enable_think) |
| text += "\n" |
| return tok(text, return_tensors="pt", return_dict=True, |
| add_special_tokens=False) |
|
|
|
|
| def err(message, etype="invalid_request_error", code=400): |
| return jsonify({"error": {"message": message, "type": etype, |
| "param": None, "code": code}}), code |
|
|
|
|
| def auth_ok(): |
| h = request.headers.get("Authorization", "") |
| if h == f"Bearer {API_KEY}": |
| return True |
| if request.headers.get("x-api-key") == API_KEY: |
| return True |
| return False |
|
|
|
|
| def gen_once(msgs, temperature, top_p, max_new, enable_think, stream_hints=None): |
| """单次生成, 返回 (text, n_in, n_out, stopped, loop, steps).""" |
| tok, model = MODEL["tok"], MODEL["m"] |
| enc = build_inputs(msgs, enable_think) |
| n_in = enc["input_ids"].shape[1] |
| if n_in + max_new > MODEL["ctx"]: |
| raise ValueError( |
| f"prompt {n_in} + max_tokens {max_new} > context {MODEL['ctx']}") |
| enc = {k: v.to(model.device) for k, v in enc.items()} |
| do_sample = temperature and temperature > 0 |
| gen_kw = dict(max_new_tokens=max_new, do_sample=do_sample, |
| repetition_penalty=1.05, pad_token_id=2, eos_token_id=EOS_IDS, |
| stopping_criteria=StoppingCriteriaList([AntiLoop()])) |
| if do_sample: |
| gen_kw["temperature"] = float(temperature) |
| gen_kw["top_p"] = float(top_p) |
| n_log = len(model._ponder_log) |
| t0 = time.time() |
| with torch.no_grad(): |
| out = model.generate(**enc, **gen_kw) |
| el = round(time.time() - t0, 1) |
| text = tok.decode(out[0][n_in:], skip_special_tokens=False) |
| entries = model._ponder_log[n_log:] |
| steps = None |
| if entries: |
| steps = round(sum(e.get("steps_mean", e.get("executed", 0)) or 0 |
| for e in entries) / len(entries), 2) |
| n_out = int(out.shape[1] - n_in) |
| stopped = int(out[0][-1]) in EOS_IDS |
| loop = (not stopped) and (n_out < max_new) |
| if stream_hints is not None: |
| stream_hints.update(elapsed=el, steps=steps) |
| return text, n_in, n_out, stopped, loop, steps |
|
|
|
|
| def to_openai(cid, model_name, reply, think, finish, n_in, n_out, steps=None, |
| elapsed=None, stream=False): |
| msg = {"role": "assistant", "content": reply} |
| if think: |
| msg["reasoning_content"] = think |
| body = { |
| "id": f"chatcmpl-{cid}", "object": "chat.completion", |
| "created": int(time.time()), "model": model_name, |
| "choices": [{"index": 0, |
| "message": msg, |
| "finish_reason": finish}], |
| "usage": {"prompt_tokens": n_in, "completion_tokens": n_out, |
| "total_tokens": n_in + n_out}, |
| } |
| if steps is not None: |
| body["samai_ponder_steps"] = steps |
| if elapsed is not None: |
| body["samai_elapsed_s"] = elapsed |
| return body |
|
|
|
|
| @app.route("/v1/models") |
| def list_models(): |
| if not auth_ok(): |
| return err("Invalid API key", code=401) |
| return jsonify({"object": "list", "data": [ |
| {"id": MODEL_NAME, "object": "model", "created": int(STATE["t0"]), |
| "owned_by": "samai"}]}) |
|
|
|
|
| @app.route("/v1/chat/completions", methods=["POST"]) |
| def chat_completions(): |
| if not auth_ok(): |
| return err("Invalid API key", code=401) |
| if not STATE["loaded"]: |
| return err("model loading", "server_error", 503) |
| d = request.get_json(force=True) |
| msgs_in = d.get("messages") or [] |
| msgs = [{"role": m.get("role", "user"), "content": str(m.get("content", ""))} |
| for m in msgs_in |
| if m.get("role") in ("system", "user", "assistant") and |
| m.get("content") is not None] |
| if not msgs: |
| return err("messages is empty") |
| enable_think = bool(d.get("enable_thinking", True)) |
| temperature = d.get("temperature", 0.6 if enable_think else 1.0) |
| temperature = 0.0 if temperature is None else float(temperature) |
| top_p = float(d.get("top_p", 0.95)) |
| max_tokens = int(d.get("max_tokens") or DEFAULT_MAX_TOKENS) |
| max_tokens = max(1, min(HARD_MAX_TOKENS, max_tokens)) |
| stream = bool(d.get("stream", False)) |
| model_name = d.get("model") or MODEL_NAME |
|
|
| with LOCK: |
| try: |
| hints = {} |
| text, n_in, n_out, stopped, loop, steps = gen_once( |
| msgs, temperature, top_p, max_tokens, enable_think, hints) |
| except ValueError as e: |
| return err(str(e), code=400) |
| except Exception as e: |
| import traceback |
| traceback.print_exc() |
| return err(repr(e)[:300], "internal_error", 500) |
| think, reply = split_think(text) |
| finish = "stop" if stopped else ("length" if not loop else "stop") |
| cid = uuid.uuid4().hex[:24] |
|
|
| if not stream: |
| body = to_openai(cid, model_name, reply, think, finish, |
| n_in, n_out, steps, hints.get("elapsed")) |
| return jsonify(body) |
|
|
| |
| def sse(): |
| def chunk(delta, fr=None): |
| c = {"id": f"chatcmpl-{cid}", "object": "chat.completion.chunk", |
| "created": int(time.time()), "model": model_name, |
| "choices": [{"index": 0, "delta": delta, "finish_reason": fr}]} |
| return f"data: {json.dumps(c, ensure_ascii=False)}\n\n" |
|
|
| yield chunk({"role": "assistant"}) |
| if think: |
| yield chunk({"reasoning_content": think}) |
| step = 8 |
| for i in range(0, len(reply), step): |
| yield chunk({"content": reply[i:i + step]}) |
| last = {"content": ""} |
| if steps is not None: |
| last["ponder_steps"] = steps |
| yield chunk(last, fr=finish) |
| u = {"prompt_tokens": n_in, "completion_tokens": n_out, |
| "total_tokens": n_in + n_out} |
| yield ("data: " + json.dumps( |
| {"id": f"chatcmpl-{cid}", "object": "chat.completion.chunk", |
| "created": int(time.time()), "model": model_name, |
| "choices": [], "usage": u}, ensure_ascii=False) + "\n\n") |
| yield "data: [DONE]\n\n" |
|
|
| return Response(sse(), mimetype="text/event-stream", |
| headers={"Cache-Control": "no-cache", |
| "X-Accel-Buffering": "no"}) |
|
|
|
|
| @app.route("/health") |
| def health(): |
| return jsonify({"loaded": STATE["loaded"], "error": STATE["error"], |
| "service": "samai-4b-openai-api", "port": PORT, |
| "ctx": MODEL["ctx"], |
| "uptime_s": round(time.time() - STATE["t0"])}) |
|
|
|
|
| @app.route("/") |
| def index(): |
| return jsonify({"service": "samai-4b OpenAI-compatible API", |
| "endpoints": ["/v1/models", "/v1/chat/completions", |
| "/health"], |
| "auth": "Authorization: Bearer <key>"}) |
|
|
|
|
| if __name__ == "__main__": |
| threading.Thread(target=build, daemon=True).start() |
| app.run(host="0.0.0.0", port=PORT, threaded=True) |
|
|
