File size: 5,794 Bytes
3c1b1a5 | 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 | #!/usr/bin/env python3
"""server.py - HTTP server for heartly-qwen-code-v3.
Serves a browser chat UI at / plus JSON endpoints /health and /chat.
The model's <decide>/<verify>/<stop>/<thinking> scaffolding is stripped by
reply_formatter (same-directory module) before the answer reaches the user.
"""
from __future__ import annotations
import argparse
import os
import threading
from typing import Optional
import torch
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from reply_formatter import format_reply
class ChatRequest(BaseModel):
prompt: str
max_new_tokens: int = 512
temperature: float = 0.7
top_p: float = 0.9
do_sample: bool = True
mode: str = "chat" # chat | debug | raw
class _State:
model: Optional[object] = None
tokenizer: Optional[object] = None
model_name: str = os.environ.get("HEARTLY_MODEL", "eivintobias/heartly-qwen-code")
lock: threading.Lock = threading.Lock()
def _format_prompt(prompt: str) -> str:
return f"User: {prompt}\nAssistant: "
def _load() -> None:
if _State.model is not None:
return
with _State.lock:
if _State.model is None:
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
device = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(_State.model_name)
model = AutoModelForCausalLM.from_pretrained(_State.model_name, torch_dtype=dtype, device_map=device)
model.eval()
_State.tokenizer, _State.model = tok, model
app = FastAPI(title="Heartly Qwen-Code v3", version="3.0")
CHAT_HTML = """<!doctype html>
<html lang="en">
<head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>Heartly Qwen-Code v3</title>
<style>
html,body{margin:0;height:100%;font-family:-apple-system,BlinkMacSystemFont,Segoe UI,Roboto,Helvetica,Arial,sans-serif;background:#0b0f14;color:#e6e8ec}
#wrap{max-width:820px;height:100vh;margin:0 auto;display:flex;flex-direction:column}
#msgs{flex:1;overflow-y:auto;padding:18px 16px 14px;display:flex;flex-direction:column;gap:12px}
.msg{max-width:82%;white-space:pre-wrap;line-height:1.45;padding:10px 14px;border-radius:14px;font-size:14px}
.user{background:#1f2430;margin-left:auto;border-radius:16px 4px 16px 16px}
.bot{background:#171b23;margin-right:auto;border-radius:4px 16px 16px 16px}
.bot.placeholder{opacity:.55}
#form{display:flex;gap:8px;padding:12px;background:#0f131a;border-top:1px solid #1d222c}
#input{flex:1;background:#171b23;border:1px solid #2a2f3c;border-radius:10px;color:#e6e8ec;padding:10px 12px;font-size:14px;outline:none}
#input::placeholder{color:#7a8190}
button{background:#2f6fec;border:none;color:#fff;border-radius:10px;padding:10px 16px;cursor:pointer;font-size:14px}
button:disabled{opacity:.5;cursor:not-allowed}
</style></head><body>
<div id="wrap"><div id="msgs"></div>
<form id="form" autocomplete="off">
<input id="input" placeholder="Ask the Heartly Qwen-Code v3 model..." autofocus>
<button id="send">Send</button>
</form></div>
<script>
const msgs=document.getElementById('msgs'),form=document.getElementById('form'),input=document.getElementById('input'),btn=document.getElementById('send');
function addMsg(c,t,ph){const d=document.createElement('div');d.className='msg '+c;if(ph)d.classList.add('placeholder');d.textContent=t;msgs.appendChild(d);msgs.scrollTop=msgs.scrollHeight;return d;}
function busy(b){btn.disabled=b;input.disabled=b;}
form.onsubmit=function(e){e.preventDefault();const p=input.value.trim();if(!p||btn.disabled)return;
addMsg('user',p);const bot=addMsg('bot','thinking...',true);busy(true);input.value='';
fetch('/chat',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({prompt:p,max_new_tokens:256,temperature:0.3,top_p:0.9,mode:'chat'})})
.then(r=>r.json()).then(d=>{bot.classList.remove('placeholder');bot.textContent=d.reply||'(no reply)';})
.catch(err=>{bot.classList.remove('placeholder');bot.textContent='error: '+err;}).finally(()=>{busy(false);input.focus();});};
</script></body></html>
"""
@app.get("/", response_class=HTMLResponse)
async def chat_page():
"""Browser chat UI - open the tab and start typing."""
return CHAT_HTML
@app.get("/health")
@torch.no_grad()
async def health():
return {"status": "ready" if _State.model is not None else "loading (loads on first /chat)", "model": _State.model_name}
@app.post("/chat")
@torch.no_grad()
async def chat(req: ChatRequest):
_load()
model, tok = _State.model, _State.tokenizer
device = next(model.parameters()).device
ids = tok.encode(_format_prompt(req.prompt), return_tensors="pt").to(device)
out = model.generate(ids, max_new_tokens=req.max_new_tokens, temperature=req.temperature, top_p=req.top_p, do_sample=req.do_sample, pad_token_id=tok.eos_token_id)
raw = tok.decode(out[0][ids.shape[1]:], skip_special_tokens=False)
reply = format_reply(raw, mode=req.mode)
return {"model": _State.model_name, "raw": raw, "reply": reply}
def main():
p = argparse.ArgumentParser(description="Heartly Qwen-Code v3 HTTP server")
p.add_argument("--model", default=os.environ.get("HEARTLY_MODEL", "eivintobias/heartly-qwen-code"))
p.add_argument("--host", default="127.0.0.1")
p.add_argument("--port", type=int, default=8000)
a = p.parse_args()
_State.model_name = a.model
import uvicorn
uvicorn.run("server:app", host=a.host, port=a.port)
if __name__ == "__main__":
main()
|