dx
Set fallback chat template on both processor and tokenizer to resolve apply_chat_template error
a2487bc | import base64 | |
| import json | |
| import time | |
| from typing import List, Optional | |
| from pydantic import BaseModel | |
| import httpx | |
| from fastapi import FastAPI, Request, HTTPException, Form, UploadFile, File | |
| from fastapi.responses import HTMLResponse, StreamingResponse | |
| from fastapi.middleware.cors import CORSMiddleware | |
| import io | |
| import threading | |
| from threading import Thread | |
| from PIL import Image | |
| import torch | |
| try: | |
| from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer | |
| except Exception as e: | |
| print("FATAL ERROR: Failed to import transformers classes.") | |
| import traceback | |
| traceback.print_exc() | |
| raise e | |
| app = FastAPI( | |
| title="Gemma 4 E2B API Space", | |
| description="FastAPI wrapper for Gemma 4 E2B via native HF Transformers.", | |
| version="1.0.0" | |
| ) | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_credentials=True, | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| MODEL_NAME = "google/gemma-4-E2B" | |
| # Global references for model and processor | |
| processor = None | |
| model = None | |
| model_loaded = False | |
| load_error = None | |
| loading_lock = threading.Lock() | |
| def load_model_if_needed(): | |
| global processor, model, model_loaded, load_error | |
| with loading_lock: | |
| if model_loaded: | |
| return | |
| try: | |
| print(f"Loading processor and model: {MODEL_NAME}...") | |
| processor = AutoProcessor.from_pretrained(MODEL_NAME) | |
| # Setup fallback chat template if the official template file is missing/uncached | |
| tokenizer = None | |
| if hasattr(processor, "tokenizer"): | |
| tokenizer = processor.tokenizer | |
| elif hasattr(processor, "image_processor") and hasattr(processor, "tokenizer"): | |
| tokenizer = processor.tokenizer | |
| has_template = False | |
| if hasattr(processor, "chat_template") and processor.chat_template: | |
| has_template = True | |
| elif tokenizer and hasattr(tokenizer, "chat_template") and tokenizer.chat_template: | |
| has_template = True | |
| if not has_template: | |
| print("Setting fallback chat template for Gemma...") | |
| fallback_template = ( | |
| "{{ bos_token }}" | |
| "{% for message in messages %}" | |
| "<start_of_turn>{{ message['role'] }}\n" | |
| "{% if message['content'] is string %}" | |
| "{{ message['content'] }}" | |
| "{% else %}" | |
| "{% for part in message['content'] %}" | |
| "{% if part['type'] == 'text' %}" | |
| "{{ part['text'] }}" | |
| "{% elif part['type'] == 'image' %}" | |
| "<image>" | |
| "{% endif %}" | |
| "{% endfor %}" | |
| "{% endif %}" | |
| "<end_of_turn>\n" | |
| "{% endfor %}" | |
| "{% if add_generation_prompt %}" | |
| "<start_of_turn>model\n" | |
| "{% endif %}" | |
| ) | |
| try: | |
| processor.chat_template = fallback_template | |
| except Exception: | |
| pass | |
| if tokenizer: | |
| tokenizer.chat_template = fallback_template | |
| try: | |
| print(f"Attempting to load model in bfloat16: {MODEL_NAME}...") | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| MODEL_NAME, | |
| device_map="cpu", | |
| torch_dtype=torch.bfloat16, | |
| low_cpu_mem_usage=True | |
| ) | |
| print("Model loaded in bfloat16 successfully!") | |
| except Exception as e_bf16: | |
| print(f"bfloat16 load failed ({e_bf16}). Falling back to float32...") | |
| model = AutoModelForImageTextToText.from_pretrained( | |
| MODEL_NAME, | |
| device_map="cpu", | |
| torch_dtype=torch.float32, | |
| low_cpu_mem_usage=True | |
| ) | |
| print("Model loaded in float32 successfully!") | |
| model_loaded = True | |
| print("Model loaded successfully!") | |
| except Exception as e: | |
| load_error = str(e) | |
| print(f"Error loading model: {load_error}") | |
| raise e | |
| def background_load_model(): | |
| try: | |
| load_model_if_needed() | |
| except Exception: | |
| pass | |
| async def startup_event(): | |
| # Start loading the model in a background thread to prevent HF Spaces boot timeout | |
| thread = threading.Thread(target=background_load_model, daemon=True) | |
| thread.start() | |
| class GenerateRequest(BaseModel): | |
| prompt: str | |
| system: Optional[str] = None | |
| temperature: Optional[float] = 0.7 | |
| class ChatMessage(BaseModel): | |
| role: str | |
| content: str | |
| class ChatRequest(BaseModel): | |
| messages: List[ChatMessage] | |
| temperature: Optional[float] = 0.7 | |
| # --------------------------------------------------------------------------- | |
| # HTML UI | |
| # --------------------------------------------------------------------------- | |
| HTML_PAGE = """<!DOCTYPE html> | |
| <html lang="es"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <title>Gemma 4 E2B — Live Console</title> | |
| <link rel="preconnect" href="https://fonts.googleapis.com"> | |
| <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin> | |
| <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800&family=JetBrains+Mono:wght@400;500;700&display=swap" rel="stylesheet"> | |
| <style> | |
| *{box-sizing:border-box;margin:0;padding:0} | |
| :root{ | |
| --bg:#0a0e17;--surface:#111827;--surface2:#1a2236; | |
| --border:rgba(255,255,255,.07);--border2:rgba(255,255,255,.12); | |
| --accent:#8b5cf6;--accent2:#a78bfa;--cyan:#22d3ee;--green:#34d399;--red:#f87171;--amber:#fbbf24; | |
| --text:#f1f5f9;--muted:#94a3b8; | |
| } | |
| html,body{height:100%;font-family:'Inter',sans-serif;background:var(--bg);color:var(--text);overflow:hidden} | |
| /* layout */ | |
| .app{display:flex;height:100vh;gap:0} | |
| .panel{display:flex;flex-direction:column;overflow:hidden} | |
| .chat-panel{flex:1;min-width:0;border-right:1px solid var(--border)} | |
| .log-panel{width:420px;flex-shrink:0;background:var(--surface)} | |
| /* header */ | |
| .hdr{padding:14px 20px;display:flex;align-items:center;gap:10px;border-bottom:1px solid var(--border);background:var(--surface)} | |
| .hdr-title{font-weight:700;font-size:1.05rem} | |
| .hdr-model{font-family:'JetBrains Mono',monospace;font-size:.78rem;color:var(--cyan);background:rgba(34,211,238,.08);padding:3px 10px;border-radius:6px;border:1px solid rgba(34,211,238,.15)} | |
| .dot{width:8px;height:8px;border-radius:50%;background:var(--green);box-shadow:0 0 8px var(--green);animation:blink 2s infinite} | |
| @keyframes blink{0%,100%{opacity:.4}50%{opacity:1}} | |
| .spacer{flex:1} | |
| .hdr-badge{font-size:.72rem;color:var(--muted);font-weight:500} | |
| /* chat area */ | |
| .chat-body{flex:1;overflow-y:auto;padding:20px;display:flex;flex-direction:column;gap:14px} | |
| .msg{max-width:82%;padding:12px 16px;border-radius:16px;font-size:.92rem;line-height:1.55;word-wrap:break-word;animation:fadeUp .25s ease} | |
| @keyframes fadeUp{from{opacity:0;transform:translateY(6px)}to{opacity:1;transform:translateY(0)}} | |
| .msg.user{align-self:flex-end;background:linear-gradient(135deg,var(--accent),#7c3aed);color:#fff;border-bottom-right-radius:4px} | |
| .msg.ai{align-self:flex-start;background:var(--surface2);border:1px solid var(--border2);border-bottom-left-radius:4px} | |
| .msg.ai pre{font-family:'JetBrains Mono',monospace;font-size:.82rem;white-space:pre-wrap;margin:0} | |
| .msg.err{align-self:center;background:rgba(248,113,113,.08);border:1px solid rgba(248,113,113,.2);color:var(--red);font-size:.82rem;max-width:90%;text-align:center} | |
| .typing{display:flex;gap:5px;padding:4px 0} | |
| .typing span{width:6px;height:6px;background:var(--muted);border-radius:50%;animation:bounce 1.4s infinite ease-in-out both} | |
| .typing span:nth-child(1){animation-delay:-.32s} | |
| .typing span:nth-child(2){animation-delay:-.16s} | |
| @keyframes bounce{0%,80%,100%{transform:scale(0)}40%{transform:scale(1)}} | |
| /* input */ | |
| .input-bar{padding:14px 20px;border-top:1px solid var(--border);display:flex;gap:10px;align-items:flex-end;background:var(--surface)} | |
| .input-bar textarea{flex:1;resize:none;background:#0d1220;border:1px solid var(--border2);border-radius:12px;padding:12px 14px;color:var(--text);font-family:inherit;font-size:.92rem;outline:none;height:46px;max-height:120px;transition:border-color .2s} | |
| .input-bar textarea:focus{border-color:rgba(139,92,246,.45)} | |
| .send-btn{width:46px;height:46px;border-radius:12px;border:none;background:linear-gradient(135deg,var(--accent),#7c3aed);color:#fff;cursor:pointer;display:flex;align-items:center;justify-content:center;transition:transform .15s,filter .15s;box-shadow:0 4px 12px rgba(139,92,246,.25)} | |
| .send-btn:hover{transform:translateY(-1px);filter:brightness(1.12)} | |
| .send-btn:active{transform:translateY(0)} | |
| .send-btn svg{width:18px;height:18px;fill:#fff} | |
| /* log panel */ | |
| .log-hdr{padding:14px 16px;border-bottom:1px solid var(--border);display:flex;align-items:center;gap:8px;font-weight:600;font-size:.9rem;background:rgba(0,0,0,.25)} | |
| .log-hdr .ico{font-size:1.1rem} | |
| .log-body{flex:1;overflow-y:auto;padding:10px;font-family:'JetBrains Mono',monospace;font-size:.76rem;line-height:1.65;color:var(--muted)} | |
| .log-entry{padding:8px 10px;border-bottom:1px solid rgba(255,255,255,.03);animation:fadeUp .2s ease} | |
| .log-ts{color:#64748b;margin-right:6px} | |
| .log-tag{font-weight:700;margin-right:4px} | |
| .log-tag.req{color:var(--amber)} | |
| .log-tag.res{color:var(--green)} | |
| .log-tag.err{color:var(--red)} | |
| .log-tag.sys{color:var(--cyan)} | |
| .log-data{color:#cbd5e1;word-break:break-all} | |
| .log-clear{margin-left:auto;background:transparent;border:1px solid var(--border2);color:var(--muted);padding:4px 10px;border-radius:6px;font-size:.7rem;cursor:pointer;font-family:inherit} | |
| .log-clear:hover{background:rgba(255,255,255,.04);color:var(--text)} | |
| /* scrollbar */ | |
| ::-webkit-scrollbar{width:6px} | |
| ::-webkit-scrollbar-track{background:transparent} | |
| ::-webkit-scrollbar-thumb{background:rgba(255,255,255,.1);border-radius:3px} | |
| ::-webkit-scrollbar-thumb:hover{background:rgba(255,255,255,.18)} | |
| @media(max-width:800px){ | |
| .app{flex-direction:column} | |
| .chat-panel{border-right:none;border-bottom:1px solid var(--border)} | |
| .log-panel{width:100%;height:220px} | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <div class="app"> | |
| <!-- CHAT --> | |
| <div class="panel chat-panel"> | |
| <div class="hdr"> | |
| <span class="dot"></span> | |
| <span class="hdr-title">Gemma 4 E2B</span> | |
| <span class="hdr-model">gemma4:e2b</span> | |
| <span class="spacer"></span> | |
| <span class="hdr-badge">CPU · Ollama · FastAPI</span> | |
| </div> | |
| <div class="chat-body" id="chat"></div> | |
| <div class="input-bar"> | |
| <textarea id="inp" rows="1" placeholder="Escribe un mensaje…"></textarea> | |
| <button class="send-btn" id="send"><svg viewBox="0 0 24 24"><path d="M2,21L23,12L2,3V10L17,12L2,14V21Z"/></svg></button> | |
| </div> | |
| </div> | |
| <!-- LOG --> | |
| <div class="panel log-panel"> | |
| <div class="log-hdr"> | |
| <span class="ico">📋</span> Request / Response Log | |
| <button class="log-clear" onclick="clearLog()">Clear</button> | |
| </div> | |
| <div class="log-body" id="log"></div> | |
| </div> | |
| </div> | |
| <script> | |
| const chat = document.getElementById('chat'); | |
| const log = document.getElementById('log'); | |
| const inp = document.getElementById('inp'); | |
| const sendBtn = document.getElementById('send'); | |
| // auto-grow textarea | |
| inp.addEventListener('input', () => { | |
| inp.style.height = '46px'; | |
| inp.style.height = inp.scrollHeight + 'px'; | |
| }); | |
| inp.addEventListener('keydown', e => { | |
| if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); send(); } | |
| }); | |
| sendBtn.addEventListener('click', send); | |
| function ts() { | |
| return new Date().toLocaleTimeString('es-MX', {hour:'2-digit',minute:'2-digit',second:'2-digit'}); | |
| } | |
| function addLog(tag, cls, text) { | |
| const d = document.createElement('div'); | |
| d.className = 'log-entry'; | |
| d.innerHTML = `<span class="log-ts">[${ts()}]</span><span class="log-tag ${cls}">${tag}</span><span class="log-data">${text}</span>`; | |
| log.appendChild(d); | |
| log.scrollTop = log.scrollHeight; | |
| } | |
| function clearLog() { log.innerHTML = ''; addLog('SYS','sys','Log cleared'); } | |
| function addMsg(role, html) { | |
| const d = document.createElement('div'); | |
| d.className = 'msg ' + role; | |
| d.innerHTML = html; | |
| chat.appendChild(d); | |
| chat.scrollTop = chat.scrollHeight; | |
| return d; | |
| } | |
| // Startup log | |
| addLog('SYS','sys','UI loaded — ready to chat with gemma4:e2b'); | |
| async function send() { | |
| const text = inp.value.trim(); | |
| if (!text) return; | |
| inp.value = ''; inp.style.height = '46px'; | |
| addMsg('user', esc(text)); | |
| addLog('REQ','req',`POST /chat prompt="${shorten(text, 120)}"`); | |
| const aiDiv = addMsg('ai', '<div class="typing"><span></span><span></span><span></span></div>'); | |
| const t0 = performance.now(); | |
| const body = new FormData(); | |
| body.append('prompt', text); | |
| try { | |
| const res = await fetch('/chat', { method: 'POST', body }); | |
| addLog('SYS','sys',`HTTP ${res.status} — streaming started`); | |
| if (!res.ok) { | |
| const errTxt = await res.text(); | |
| aiDiv.className = 'msg err'; | |
| aiDiv.textContent = 'Error: ' + errTxt; | |
| addLog('ERR','err', errTxt); | |
| return; | |
| } | |
| const reader = res.body.getReader(); | |
| const dec = new TextDecoder(); | |
| let full = '', buf = '', tokens = 0; | |
| while (true) { | |
| const { done, value } = await reader.read(); | |
| if (done) break; | |
| buf += dec.decode(value, { stream: true }); | |
| const lines = buf.split('\\n'); | |
| buf = lines.pop(); | |
| for (const ln of lines) { | |
| const l = ln.trim(); | |
| if (l.startsWith('data: ')) { | |
| let chunk = l.substring(6).replace(/\\\\n/g, '\\n'); | |
| if (chunk === '[PING]' || chunk === '[ERROR]') continue; | |
| full += chunk; | |
| tokens++; | |
| } | |
| } | |
| aiDiv.innerHTML = '<pre>' + esc(full) + '</pre>'; | |
| chat.scrollTop = chat.scrollHeight; | |
| } | |
| const elapsed = ((performance.now() - t0) / 1000).toFixed(1); | |
| const tps = (tokens / (parseFloat(elapsed) || 1)).toFixed(1); | |
| addLog('RES','res',`${tokens} tokens in ${elapsed}s (${tps} tok/s)`); | |
| addLog('RES','res',`Response: "${shorten(full, 200)}"`); | |
| if (!full.trim()) { | |
| aiDiv.innerHTML = '<em style="color:var(--muted)">(sin respuesta)</em>'; | |
| addLog('SYS','sys','Empty response from model'); | |
| } | |
| } catch (err) { | |
| aiDiv.className = 'msg err'; | |
| aiDiv.textContent = 'Connection error: ' + err.message; | |
| addLog('ERR','err', err.message); | |
| } | |
| } | |
| function esc(s) { return s.replace(/&/g,'&').replace(/</g,'<').replace(/>/g,'>'); } | |
| function shorten(s, n) { return s.length > n ? s.slice(0, n) + '…' : s; } | |
| </script> | |
| </body> | |
| </html>""" | |
| # --------------------------------------------------------------------------- | |
| # Routes | |
| # --------------------------------------------------------------------------- | |
| # --------------------------------------------------------------------------- | |
| # Helper function to extract PIL Images and filter messages list | |
| # --------------------------------------------------------------------------- | |
| def extract_images_and_clean_messages(messages): | |
| images = [] | |
| cleaned_messages = [] | |
| for msg in messages: | |
| role = msg.get("role") | |
| content = msg.get("content") | |
| if isinstance(content, list): | |
| cleaned_parts = [] | |
| for part in content: | |
| if part.get("type") == "image": | |
| img = part.get("image") | |
| if img: | |
| images.append(img) | |
| # We keep the dict structure but remove raw PIL objects for serialization, | |
| # or keep it if processor needs it. Standard transformers template | |
| # accepts type="image" and ignores or expects it. | |
| cleaned_parts.append(part) | |
| else: | |
| cleaned_parts.append(part) | |
| cleaned_messages.append({"role": role, "content": cleaned_parts}) | |
| else: | |
| cleaned_messages.append({"role": role, "content": content}) | |
| return cleaned_messages, images | |
| # --------------------------------------------------------------------------- | |
| # Helper function for streaming generation | |
| # --------------------------------------------------------------------------- | |
| async def run_generation(messages): | |
| global model_loaded, load_error | |
| if not model_loaded: | |
| if load_error: | |
| yield f"data: [ERROR] Model load failed: {load_error}\n\n" | |
| else: | |
| yield "data: [ERROR] Model is still loading. Please try again in a few seconds.\n\n" | |
| return | |
| try: | |
| # Extract images and clean messages | |
| cleaned_messages, images = extract_images_and_clean_messages(messages) | |
| # Apply chat template (yields text prompt) | |
| text = processor.apply_chat_template( | |
| cleaned_messages, | |
| tokenize=False, | |
| add_generation_prompt=True | |
| ) | |
| # Build processor inputs | |
| if images: | |
| inputs = processor(text=text, images=images, return_tensors="pt") | |
| else: | |
| inputs = processor(text=text, return_tensors="pt") | |
| # Move to CPU / model device | |
| inputs = {k: v.to(model.device) for k, v in inputs.items()} | |
| streamer = TextIteratorStreamer( | |
| processor.tokenizer, | |
| skip_prompt=True, | |
| skip_special_tokens=True | |
| ) | |
| generate_kwargs = dict( | |
| **inputs, | |
| streamer=streamer, | |
| max_new_tokens=1024, | |
| do_sample=True, | |
| temperature=0.7 | |
| ) | |
| # Run model generate in a daemon thread so it does not block the event loop | |
| thread = Thread(target=model.generate, kwargs=generate_kwargs, daemon=True) | |
| thread.start() | |
| for new_text in streamer: | |
| if new_text: | |
| text_escaped = new_text.replace("\n", "\\n") | |
| yield f"data: {text_escaped}\n\n" | |
| except Exception as e: | |
| yield f"data: [ERROR] Generation error: {str(e)}\n\n" | |
| def run_generation_sync(messages, temperature=0.7): | |
| if not model_loaded: | |
| if load_error: | |
| raise HTTPException(status_code=503, detail=f"Model load failed: {load_error}") | |
| raise HTTPException(status_code=503, detail="Model is still loading. Please try again.") | |
| cleaned_messages, images = extract_images_and_clean_messages(messages) | |
| text = processor.apply_chat_template( | |
| cleaned_messages, | |
| tokenize=False, | |
| add_generation_prompt=True | |
| ) | |
| if images: | |
| inputs = processor(text=text, images=images, return_tensors="pt") | |
| else: | |
| inputs = processor(text=text, return_tensors="pt") | |
| inputs = {k: v.to(model.device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=1024, | |
| do_sample=True, | |
| temperature=temperature | |
| ) | |
| input_len = inputs["input_ids"].shape[1] | |
| generated_tokens = outputs[0][input_len:] | |
| return processor.decode(generated_tokens, skip_special_tokens=True) | |
| # --------------------------------------------------------------------------- | |
| # Routes | |
| # --------------------------------------------------------------------------- | |
| def root(request: Request): | |
| accept = request.headers.get("accept", "") | |
| if "text/html" in accept: | |
| return HTMLResponse(content=HTML_PAGE) | |
| return { | |
| "status": "online", | |
| "model": MODEL_NAME, | |
| "message": "Gemma 4 E2B API is running natively.", | |
| "endpoints": { | |
| "chat_form": "/chat", | |
| "vision": "/vision", | |
| "generate": "/api/generate", | |
| "chat_json": "/api/chat", | |
| "health": "/health", | |
| "docs": "/docs" | |
| } | |
| } | |
| async def health_check(): | |
| global model_loaded, load_error | |
| if model_loaded: | |
| return {"status": "ok", "model": MODEL_NAME, "loaded": True} | |
| elif load_error: | |
| raise HTTPException(status_code=503, detail=f"Model failed to load: {load_error}") | |
| else: | |
| return {"status": "loading", "model": MODEL_NAME, "loaded": False} | |
| # --------------------------------------------------------------------------- | |
| # POST /chat – FormData → SSE | |
| # --------------------------------------------------------------------------- | |
| async def chat_form(prompt: str = Form(...)): | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [{"type": "text", "text": prompt}] | |
| } | |
| ] | |
| return StreamingResponse( | |
| run_generation(messages), | |
| media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"} | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # POST /vision – FormData (prompt + imagen) → SSE | |
| # --------------------------------------------------------------------------- | |
| async def vision_form(prompt: str = Form(...), imagen: UploadFile = File(...)): | |
| try: | |
| image_bytes = await imagen.read() | |
| pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") | |
| except Exception as e: | |
| raise HTTPException(status_code=400, detail=f"Invalid image file: {str(e)}") | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [ | |
| {"type": "image", "image": pil_image}, | |
| {"type": "text", "text": prompt} | |
| ] | |
| } | |
| ] | |
| return StreamingResponse( | |
| run_generation(messages), | |
| media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"} | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # JSON endpoints | |
| # --------------------------------------------------------------------------- | |
| async def generate_completion(payload: GenerateRequest): | |
| messages = [ | |
| { | |
| "role": "user", | |
| "content": [{"type": "text", "text": payload.prompt}] | |
| } | |
| ] | |
| if payload.system: | |
| messages.insert(0, { | |
| "role": "system", | |
| "content": [{"type": "text", "text": payload.system}] | |
| }) | |
| try: | |
| response_text = run_generation_sync(messages, temperature=payload.temperature) | |
| return { | |
| "model": MODEL_NAME, | |
| "response": response_text, | |
| "done": True | |
| } | |
| except HTTPException as he: | |
| raise he | |
| except Exception as e: | |
| raise HTTPException(status_code=500, detail=f"Generation failed: {str(e)}") | |
| async def chat_completion(payload: ChatRequest): | |
| messages = [] | |
| for msg in payload.messages: | |
| messages.append({ | |
| "role": msg.role, | |
| "content": [{"type": "text", "text": msg.content}] | |
| }) | |
| try: | |
| response_text = run_generation_sync(messages, temperature=payload.temperature) | |
| return { | |
| "model": MODEL_NAME, | |
| "message": { | |
| "role": "assistant", | |
| "content": response_text | |
| }, | |
| "done": True | |
| } | |
| except HTTPException as he: | |
| raise he | |
| except Exception as e: | |
| raise HTTPException(status_code=503, detail=f"Generation failed: {str(e)}") | |