Spaces:
Running on Zero
Running on Zero
| import os | |
| import re | |
| import time | |
| import traceback | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Tuple | |
| from config import * | |
| import torch | |
| from fastapi import Request | |
| from fastapi.responses import HTMLResponse, JSONResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from gradio import Server | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| # ── ZeroGPU Detection ─────────────────────────────────────────────────── | |
| IS_ZEROGPU = False | |
| try: | |
| import spaces | |
| from spaces.config import Config as _ZeroGPUConfig | |
| IS_ZEROGPU = bool(_ZeroGPUConfig.zero_gpu) | |
| except ImportError: | |
| class _DummySpaces: | |
| def GPU(fn): | |
| return fn | |
| spaces = _DummySpaces() | |
| # ── Configuration ─────────────────────────────────────────────────────── | |
| DEFAULT_SYSTEM_PROMPT = ( | |
| "You are ChatUI, a helpful local AI assistant. " | |
| "You provide accurate, detailed, and thoughtful responses." | |
| ) | |
| STATIC_DIR = Path(__file__).parent / "static" | |
| os.makedirs(STATIC_DIR, exist_ok=True) | |
| # ── Fallback chat templates for tokenizers missing .chat_template ─────── | |
| _FALLBACK_TEMPLATES = { | |
| "chatml": ( | |
| "{% for message in messages %}" | |
| "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}" | |
| "{% endfor %}" | |
| "{% if add_generation_prompt %}" | |
| "{{'<|im_start|>assistant\n'}}" | |
| "{% endif %}" | |
| ), | |
| "llama2": ( | |
| "{% if messages[0]['role'] == 'system' %}" | |
| "{{'<<SYS>>\n' + messages[0]['content'] + '\n<</SYS>>\n\n'}}" | |
| "{% set loop_messages = messages[1:] %}" | |
| "{% else %}" | |
| "{% set loop_messages = messages %}" | |
| "{% endif %}" | |
| "{% for message in loop_messages %}" | |
| "{% if message['role'] == 'user' %}" | |
| "{{'[INST] ' + message['content'] + ' [/INST]'}}" | |
| "{% elif message['role'] == 'assistant' %}" | |
| "{{message['content'] + '</s>'}}" | |
| "{% endif %}" | |
| "{% endfor %}" | |
| ), | |
| "phi3": ( | |
| "{% for message in messages %}" | |
| "{% if message['role'] == 'system' %}" | |
| "{{'<|system|>\n' + message['content'] + '<|end|>\n'}}" | |
| "{% elif message['role'] == 'user' %}" | |
| "{{'<|user|>\n' + message['content'] + '<|end|>\n'}}" | |
| "{% elif message['role'] == 'assistant' %}" | |
| "{{'<|assistant|>\n' + message['content'] + '<|end|>\n'}}" | |
| "{% endif %}" | |
| "{% endfor %}" | |
| "{% if add_generation_prompt %}" | |
| "{{'<|assistant|>\n'}}" | |
| "{% endif %}" | |
| ), | |
| "qwen": ( | |
| "{% for message in messages %}" | |
| "{% if message['role'] == 'system' %}" | |
| "{{'<|im_start|>system\n' + message['content'] + '<|im_end|>\n'}}" | |
| "{% elif message['role'] == 'user' %}" | |
| "{{'<|im_start|>user\n' + message['content'] + '<|im_end|>\n'}}" | |
| "{% elif message['role'] == 'assistant' %}" | |
| "{{'<|im_start|>assistant\n' + message['content'] + '<|im_end|>\n'}}" | |
| "{% endif %}" | |
| "{% endfor %}" | |
| "{% if add_generation_prompt %}" | |
| "{{'<|im_start|>assistant\n'}}" | |
| "{% endif %}" | |
| ), | |
| } | |
| def _pick_fallback_template(model_id: str) -> str: | |
| mid = model_id.lower() | |
| if "phi" in mid: | |
| return _FALLBACK_TEMPLATES["phi3"] | |
| if "qwen" in mid: | |
| return _FALLBACK_TEMPLATES["qwen"] | |
| if "tinyllama" in mid or "llama" in mid: | |
| return _FALLBACK_TEMPLATES["llama2"] | |
| return _FALLBACK_TEMPLATES["chatml"] | |
| def _ensure_chat_template(tokenizer, model_id: str): | |
| if getattr(tokenizer, "chat_template", None): | |
| return | |
| tpl = _pick_fallback_template(model_id) | |
| tokenizer.chat_template = tpl | |
| print(f"[chat_template] injected fallback template for {model_id}") | |
| # ── Model cache (lazy) ────────────────────────────────────────────────── | |
| _MODEL_CACHE: Dict[str, Dict[str, Any]] = {} | |
| def _device() -> str: | |
| if IS_ZEROGPU: | |
| return "cuda" | |
| if torch.cuda.is_available(): | |
| return "cuda" | |
| if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): | |
| return "mps" | |
| return "cpu" | |
| def _dtype(): | |
| if IS_ZEROGPU or torch.cuda.is_available(): | |
| return torch.float16 | |
| return torch.float32 | |
| def get_model(model_id: str) -> Dict[str, Any]: | |
| if model_id in _MODEL_CACHE: | |
| return _MODEL_CACHE[model_id] | |
| if model_id not in MODEL_REGISTRY: | |
| raise ValueError(f"Unknown model: {model_id}") | |
| device = _device() | |
| dtype = _dtype() | |
| tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) | |
| if tokenizer.pad_token_id is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| _ensure_chat_template(tokenizer, model_id) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_id, | |
| torch_dtype=dtype, | |
| trust_remote_code=True, | |
| low_cpu_mem_usage=True, | |
| ).to(device) | |
| model.eval() | |
| _MODEL_CACHE[model_id] = { | |
| "tokenizer": tokenizer, | |
| "model": model, | |
| "device": device, | |
| "loaded_at": time.time(), | |
| } | |
| return _MODEL_CACHE[model_id] | |
| # ── Generation core ───────────────────────────────────────────────────── | |
| def _flatten_content(content) -> str: | |
| if isinstance(content, str): | |
| return content | |
| parts: List[str] = [] | |
| for c in content or []: | |
| if isinstance(c, dict) and c.get("type") == "text": | |
| parts.append(c.get("text", "")) | |
| return "\n".join(parts).strip() | |
| def _build_prompt(model_id: str, messages: List[Dict[str, Any]], | |
| system_prompt: str) -> str: | |
| bundle = get_model(model_id) | |
| tokenizer = bundle["tokenizer"] | |
| flat_messages: List[Dict[str, str]] = [] | |
| if system_prompt and system_prompt.strip(): | |
| flat_messages.append({"role": "system", "content": system_prompt.strip()}) | |
| for m in messages: | |
| role = m.get("role", "user") | |
| if role == "system": | |
| continue | |
| text = _flatten_content(m.get("content", "")) | |
| if text: | |
| flat_messages.append({"role": role, "content": text}) | |
| if getattr(tokenizer, "chat_template", None): | |
| try: | |
| return tokenizer.apply_chat_template( | |
| flat_messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| except Exception as e: | |
| print(f"[chat_template] apply failed for {model_id}: {e}; using manual fallback") | |
| # Manual fallback | |
| lines = [f"{m['role']}: {m['content']}" for m in flat_messages] | |
| lines.append("assistant:") | |
| return "\n".join(lines) + "\n" | |
| def generate(model_id: str, messages: List[Dict[str, Any]], | |
| system_prompt: str, max_tokens: int, temperature: float, | |
| top_p: float = 0.95) -> Tuple[str, str]: | |
| bundle = get_model(model_id) | |
| tokenizer = bundle["tokenizer"] | |
| model = bundle["model"] | |
| device = bundle["device"] | |
| prompt = _build_prompt(model_id, messages, system_prompt) | |
| inputs = tokenizer(prompt, return_tensors="pt").to(device) | |
| do_sample = temperature > 0.01 | |
| with torch.inference_mode(): | |
| out = model.generate( | |
| **inputs, | |
| max_new_tokens=max_tokens, | |
| do_sample=do_sample, | |
| temperature=max(temperature, 1e-3) if do_sample else 1.0, | |
| top_p=top_p if do_sample else 1.0, | |
| pad_token_id=tokenizer.pad_token_id, | |
| eos_token_id=tokenizer.eos_token_id, | |
| ) | |
| new_tokens = out[0][inputs["input_ids"].shape[1]:] | |
| text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip() | |
| reasoning = "" | |
| think_match = re.search(r"<think[^>]*>(.*?)</think\s*>", text, re.DOTALL | re.IGNORECASE) | |
| if think_match: | |
| reasoning = think_match.group(1).strip() | |
| text = re.sub(r"<think[^>]*>.*?</think\s*>", "", text, flags=re.DOTALL | re.IGNORECASE).strip() | |
| return text, reasoning | |
| # ── Agent tools ───────────────────────────────────────────────────────── | |
| def _tool_calculator(expression: str) -> str: | |
| try: | |
| if not re.fullmatch(r"[\d+\-*/().%\s]+", expression): | |
| return "Error: only numeric expressions are allowed." | |
| result = eval(expression, {"__builtins__": {}}, {}) | |
| return f"= {result}" | |
| except Exception as e: | |
| return f"Error: {e}" | |
| def _tool_web_search(query: str) -> str: | |
| return ( | |
| f"[stub search] I cannot reach the web (no API key configured), " | |
| f"but for query '{query}' you might want to check Wikipedia, " | |
| f"Stack Overflow, or arXiv. Use the calculator for any maths." | |
| ) | |
| def _tool_file_read(path: str) -> str: | |
| safe_root = Path(__file__).parent.resolve() | |
| try: | |
| target = (safe_root / path).resolve() | |
| target.relative_to(safe_root) | |
| if not target.exists() or not target.is_file(): | |
| return f"Error: file not found: {path}" | |
| text = target.read_text(encoding="utf-8", errors="ignore") | |
| if len(text) > 4000: | |
| text = text[:4000] + "\n…[truncated]" | |
| return text | |
| except Exception as e: | |
| return f"Error: {e}" | |
| def _tool_now(_arg: str = "") -> str: | |
| return time.strftime("%Y-%m-%d %H:%M:%S %Z") | |
| AGENT_TOOLS = { | |
| "calculator": _tool_calculator, | |
| "web_search": _tool_web_search, | |
| "file_read": _tool_file_read, | |
| "now": _tool_now, | |
| } | |
| AGENT_TOOL_DESCRIPTION = """You have access to the following tools. To use one, | |
| reply with a single line that starts with `TOOL:` followed by the tool name and | |
| its argument in parentheses, e.g.: | |
| TOOL: calculator(2 * (3 + 4)) | |
| TOOL: web_search(latest news about mars) | |
| TOOL: file_read(static/index.html) | |
| TOOL: now() | |
| After receiving the tool output, you may use another tool or give your final | |
| answer. When you are done, start your final message with `FINAL:`. | |
| """ | |
| def run_agent(model_id: str, user_text: str, system_prompt: str, | |
| max_tokens: int, temperature: float, | |
| max_iters: int = 4) -> Dict[str, Any]: | |
| history = [ | |
| {"role": "system", "content": system_prompt + "\n\n" + AGENT_TOOL_DESCRIPTION}, | |
| {"role": "user", "content": user_text}, | |
| ] | |
| trace: List[Dict[str, str]] = [] | |
| for i in range(max_iters): | |
| answer, _reason = generate( | |
| model_id, history, | |
| system_prompt="", | |
| max_tokens=max_tokens, temperature=temperature, | |
| ) | |
| answer_stripped = answer.strip() | |
| if answer_stripped.upper().startswith("FINAL:"): | |
| final = answer_stripped[6:].strip() | |
| trace.append({"step": i + 1, "type": "final", "content": final}) | |
| return {"final": final, "trace": trace} | |
| m = re.match(r"\s*TOOL:\s*([A-Za-z_]+)\s*\((.*)\)\s*$", | |
| answer_stripped, re.DOTALL) | |
| if not m: | |
| trace.append({"step": i + 1, "type": "final", "content": answer_stripped}) | |
| return {"final": answer_stripped, "trace": trace} | |
| tool_name = m.group(1).strip().lower() | |
| tool_arg = m.group(2).strip() | |
| trace.append({ | |
| "step": i + 1, | |
| "type": "tool_call", | |
| "tool": tool_name, | |
| "argument": tool_arg, | |
| }) | |
| fn = AGENT_TOOLS.get(tool_name) | |
| if fn is None: | |
| tool_result = (f"Error: unknown tool '{tool_name}'. " | |
| f"Available: {list(AGENT_TOOLS.keys())}") | |
| else: | |
| try: | |
| tool_result = fn(tool_arg) | |
| except Exception as e: | |
| tool_result = f"Error: {e}" | |
| trace.append({ | |
| "step": i + 1, | |
| "type": "tool_result", | |
| "tool": tool_name, | |
| "output": tool_result, | |
| }) | |
| history.append({"role": "assistant", "content": answer_stripped}) | |
| history.append({ | |
| "role": "user", | |
| "content": (f"TOOL OUTPUT: {tool_result}\nNow either call another " | |
| f"tool or finish with FINAL: <your answer>."), | |
| }) | |
| return { | |
| "final": "I ran out of steps before producing a final answer.", | |
| "trace": trace, | |
| } | |
| # ── Gradio Server App ─────────────────────────────────────────────────── | |
| app = Server(title="ChatUI Studio", version="2.0.0") | |
| async def api_list_models(): | |
| return JSONResponse({ | |
| "status": "success", | |
| "default": DEFAULT_MODEL, | |
| "models": [ | |
| {"id": mid, **meta} | |
| for mid, meta in MODEL_REGISTRY.items() | |
| ], | |
| "cached": list(_MODEL_CACHE.keys()), | |
| }) | |
| async def api_list_tasks(): | |
| return JSONResponse({ | |
| "status": "success", | |
| "tasks": [ | |
| {"id": tid, **meta} | |
| for tid, meta in TASK_REGISTRY.items() | |
| ], | |
| }) | |
| async def api_load_model(req: Request): | |
| body = await req.json() | |
| model_id = body.get("model_id") or DEFAULT_MODEL | |
| try: | |
| if model_id not in MODEL_REGISTRY: | |
| return JSONResponse({"status": "error", "message": "Unknown model."}) | |
| get_model(model_id) | |
| return JSONResponse({ | |
| "status": "success", | |
| "model": model_id, | |
| "device": _MODEL_CACHE[model_id]["device"], | |
| }) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return JSONResponse({"status": "error", "message": str(e)}) | |
| async def api_run_inference(req: Request): | |
| try: | |
| body = await req.json() | |
| messages: List[Dict[str, Any]] = body.get("messages", []) | |
| model_id = body.get("model_id") or DEFAULT_MODEL | |
| task = body.get("task") or "chat" | |
| max_tokens_i = int(body.get("max_tokens", 512)) | |
| temperature_f = float(body.get("temperature", 0.7)) | |
| system_prompt = body.get("system_prompt", "") or "" | |
| if model_id not in MODEL_REGISTRY: | |
| return JSONResponse({"status": "error", "message": f"Unknown model: {model_id}"}) | |
| task_meta = TASK_REGISTRY.get(task, TASK_REGISTRY["chat"]) | |
| sys_prompt = (system_prompt.strip() or task_meta["system_prompt"]) | |
| if task == "agent": | |
| user_text = _flatten_content(messages[-1].get("content", "")) if messages else "" | |
| result = run_agent( | |
| model_id=model_id, | |
| user_text=user_text, | |
| system_prompt=sys_prompt, | |
| max_tokens=max_tokens_i, | |
| temperature=temperature_f, | |
| ) | |
| return JSONResponse({ | |
| "status": "success", | |
| "task": "agent", | |
| "model": model_id, | |
| "content": result["final"], | |
| "trace": result["trace"], | |
| "reasoning_content": "", | |
| }) | |
| content, reasoning = generate( | |
| model_id=model_id, | |
| messages=messages, | |
| system_prompt=sys_prompt, | |
| max_tokens=max_tokens_i, | |
| temperature=temperature_f, | |
| ) | |
| return JSONResponse({ | |
| "status": "success", | |
| "task": task, | |
| "model": model_id, | |
| "content": content, | |
| "reasoning_content": reasoning or "", | |
| }) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return JSONResponse({"status": "error", "message": str(e)}) | |
| async def homepage(): | |
| html_path = STATIC_DIR / "index.html" | |
| if html_path.exists(): | |
| with open(html_path, "r", encoding="utf-8") as f: | |
| return HTMLResponse(content=f.read(), status_code=200) | |
| return HTMLResponse( | |
| content="<h1>Frontend is building. Please refresh in a few seconds...</h1>", | |
| status_code=200, | |
| ) | |
| async def healthz(): | |
| return JSONResponse({ | |
| "status": "ok", | |
| "zerogpu": IS_ZEROGPU, | |
| "device": _device(), | |
| "models_cached": list(_MODEL_CACHE.keys()), | |
| "models_registered": list(MODEL_REGISTRY.keys()), | |
| }) | |
| app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") | |
| if __name__ == "__main__": | |
| import argparse | |
| import inspect | |
| parser = argparse.ArgumentParser( | |
| description="ChatUI Studio — local AI chat server", | |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, | |
| ) | |
| parser.add_argument( | |
| "--host", "--server-name", dest="server_name", default="0.0.0.0", | |
| help="Host/IP to bind to.", | |
| ) | |
| parser.add_argument( | |
| "--port", "--server-port", dest="server_port", type=int, default=7860, | |
| help="Port to listen on.", | |
| ) | |
| parser.add_argument( | |
| "--share", dest="share", action="store_true", | |
| help="Create a public Gradio share link.", | |
| ) | |
| parser.add_argument( | |
| "--no-share", dest="share", action="store_false", | |
| help="Disable public share link.", | |
| ) | |
| parser.set_defaults(share=False) | |
| parser.add_argument( | |
| "--show-error", dest="show_error", action="store_true", | |
| help="Show full error tracebacks in the UI.", | |
| ) | |
| parser.add_argument( | |
| "--no-show-error", dest="show_error", action="store_false", | |
| help="Hide error tracebacks in the UI.", | |
| ) | |
| parser.set_defaults(show_error=True) | |
| parser.add_argument( | |
| "--quiet", action="store_true", | |
| help="Suppress Gradio startup messages.", | |
| ) | |
| parser.add_argument( | |
| "--auth", nargs=2, metavar=("USER", "PASS"), default=None, | |
| help="Enable basic auth with USER and PASS.", | |
| ) | |
| parser.add_argument( | |
| "--ssl-keyfile", default=None, | |
| help="Path to SSL key file (enables HTTPS).", | |
| ) | |
| parser.add_argument( | |
| "--ssl-certfile", default=None, | |
| help="Path to SSL certificate file.", | |
| ) | |
| parser.add_argument( | |
| "--ssl-keyfile-password", default=None, | |
| help="Password for the SSL key file, if any.", | |
| ) | |
| parser.add_argument( | |
| "--favicon-path", default=None, | |
| help="Path to a custom favicon.", | |
| ) | |
| parser.add_argument( | |
| "--allowed-paths", nargs="*", default=None, | |
| help="Extra paths Gradio is allowed to serve from.", | |
| ) | |
| parser.add_argument( | |
| "--blocked-paths", nargs="*", default=None, | |
| help="Paths Gradio must never serve.", | |
| ) | |
| parser.add_argument( | |
| "--max-file-size", default=None, | |
| help="Max upload size (e.g. '10mb', '1gb').", | |
| ) | |
| args = parser.parse_args() | |
| launch_kwargs = { | |
| "server_name": args.server_name, | |
| "server_port": args.server_port, | |
| "share": args.share, | |
| "show_error": args.show_error, | |
| "quiet": args.quiet, | |
| "auth": tuple(args.auth) if args.auth else None, | |
| } | |
| if args.ssl_keyfile: | |
| launch_kwargs["ssl_keyfile"] = args.ssl_keyfile | |
| if args.ssl_certfile: | |
| launch_kwargs["ssl_certfile"] = args.ssl_certfile | |
| if args.ssl_keyfile_password: | |
| launch_kwargs["ssl_keyfile_password"] = args.ssl_keyfile_password | |
| if args.favicon_path: | |
| launch_kwargs["favicon_path"] = args.favicon_path | |
| if args.allowed_paths: | |
| launch_kwargs["allowed_paths"] = args.allowed_paths | |
| if args.blocked_paths: | |
| launch_kwargs["blocked_paths"] = args.blocked_paths | |
| if args.max_file_size: | |
| launch_kwargs["max_file_size"] = args.max_file_size | |
| # Drop kwargs that aren't supported by this Gradio version. | |
| try: | |
| supported = set(inspect.signature(app.launch).parameters.keys()) | |
| if "kwargs" not in supported: | |
| launch_kwargs = {k: v for k, v in launch_kwargs.items() if k in supported} | |
| except (TypeError, ValueError): | |
| pass | |
| app.launch(**launch_kwargs) |