diffusion-chatbot

Sleeping

App Files Files Community

Andhs commited on Jul 3

Commit

2424bab

verified ·

1 Parent(s): b8d11d8

Upload app.py

Browse files

Files changed (1) hide show

app.py +338 -40

app.py CHANGED Viewed

@@ -10,6 +10,337 @@ from flask import Flask, request, jsonify, Response
 from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
 # 1. Environment Parsing & Architecture Strategy Mapping
 MODEL_NAME = os.getenv("MODEL_NAME", "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1")
 IS_DIFFUSION = "diffusion" in MODEL_NAME.lower()
@@ -79,11 +410,13 @@ def clone_past_key_values(pkv):
     """Clone KV-cache. Fast path for tuples and Cache objects; falls back to deepcopy."""
     if pkv is None:
         return None
     if isinstance(pkv, tuple):
         return tuple(
             (k.clone() if k is not None else None, v.clone() if v is not None else None)
             for k, v in pkv
         )
     if hasattr(pkv, 'key_cache') and hasattr(pkv, 'value_cache'):
         try:
             new_cache = pkv.__class__()
@@ -95,6 +428,7 @@ def clone_past_key_values(pkv):
             return new_cache
         except Exception:
             pass
     return copy.deepcopy(pkv)
@@ -332,6 +666,7 @@ def load_model():
             torch_dtype=torch.bfloat16,
             trust_remote_code=False
         ).to(device).eval()
     if not IS_DIFFUSION:
         try:
             model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
@@ -344,7 +679,7 @@ def load_model():
         MODEL_NAME,
         trust_remote_code=IS_DIFFUSION
     )
-    print("Model loaded into memory workspace.")
 @app.route('/health', methods=['GET'])
@@ -367,6 +702,7 @@ def generate_text():
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": prompt}
     ]
     encoded = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
@@ -511,7 +847,7 @@ def generate_text_sse():
             accumulated = []
             for text in streamer:
-                if not text:
                     continue
                 accumulated.append(text)
                 current = "".join(accumulated)
@@ -526,48 +862,10 @@ def generate_text_sse():
     )
-HTML_UI = """<!DOCTYPE html>
-<html lang="en">
-<head>
-<meta charset="UTF-8">
-<meta name="viewport" content="width=device-width, initial-scale=1.0">
-<title>LLM API Tester</title>
-<style>
-*{box-sizing:border-box}body{font-family:-apple-system,BlinkMacSystemFont,"Segoe UI",Roboto,sans-serif;max-width:900px;margin:0 auto;padding:20px;background:#f5f5f5;color:#333}h1{margin-top:0;font-size:1.5rem}.card{background:#fff;border-radius:8px;padding:20px;margin-bottom:16px;box-shadow:0 1px 3px rgba(0,0,0,.1)}label{display:block;font-weight:600;margin-bottom:6px;font-size:.9rem}textarea,input,select{width:100%;padding:10px;border:1px solid #ddd;border-radius:6px;font-size:.95rem;font-family:inherit}textarea{resize:vertical;min-height:80px}.row{display:grid;grid-template-columns:repeat(auto-fit,minmax(140px,1fr));gap:12px;margin-bottom:12px}.field{margin-bottom:12px}button{background:#2563eb;color:#fff;border:none;padding:10px 20px;border-radius:6px;font-size:1rem;cursor:pointer;font-weight:600}button:hover{background:#1d4ed8}button:disabled{background:#93c5fd;cursor:not-allowed}.output{background:#1e1e1e;color:#e4e4e4;padding:16px;border-radius:6px;font-family:ui-monospace,SFMono-Regular,"SF Mono",Menlo,Consolas,monospace;font-size:.9rem;white-space:pre-wrap;word-break:break-word;min-height:120px;max-height:500px;overflow-y:auto}.output:empty::before{content:"Response will appear here...";color:#666}.status{font-size:.85rem;color:#666;margin-top:8px}.error{color:#dc2626}.success{color:#16a34a}.route-badge{display:inline-block;background:#e5e7eb;padding:2px 8px;border-radius:4px;font-size:.8rem;font-weight:600;margin-bottom:8px}
-</style>
-</head>
-<body>
-<h1>LLM API Tester</h1>
-<div class="card">
-<div class="route-badge" id="routeBadge">/generate</div>
-<div class="field"><label for="route">Route</label><select id="route"><option value="/generate">/generate (sync JSON)</option><option value="/generate_stream">/generate_stream (sync JSON + intermediates)</option><option value="/generate_sse">/generate_sse (streaming SSE)</option></select></div>
-<div class="field"><label for="prompt">Prompt</label><textarea id="prompt" placeholder="Enter your prompt here...">Write a short poem about stars</textarea></div>
-<div class="row"><div class="field"><label for="max_new_tokens">max_new_tokens</label><input type="number" id="max_new_tokens" value="150" min="1" max="2048"></div><div class="field"><label for="temperature">temperature</label><input type="number" id="temperature" value="0.0" min="0" max="2" step="0.1"></div><div class="field"><label for="steps">steps (diffusion)</label><input type="number" id="steps" value="256" min="1"></div><div class="field"><label for="block_size">block_size</label><input type="number" id="block_size" value="32" min="1"></div></div>
-<div class="row"><div class="field"><label for="cfg_scale">cfg_scale</label><input type="number" id="cfg_scale" value="0.0" min="0" step="0.1"></div><div class="field"><label for="remasking">remasking</label><select id="remasking"><option value="low_confidence">low_confidence</option><option value="random">random</option></select></div><div class="field"><label for="capture_interval">capture_interval</label><input type="number" id="capture_interval" value="10" min="1"></div></div>
-<button id="sendBtn">Send Request</button>
-<div class="status" id="status"></div>
-</div>
-<div class="card"><label>Response</label><div class="output" id="output"></div></div>
-<script>
-const $=id=>document.getElementById(id);
-const routeSelect=$('route');const routeBadge=$('routeBadge');const sendBtn=$('sendBtn');const output=$('output');const status=$('status');
-routeSelect.addEventListener('change',()=>{routeBadge.textContent=routeSelect.value;});
-function setStatus(msg,isError=false){status.textContent=msg;status.className='status '+(isError?'error':'success');}
-function appendOutput(text,clear=false){if(clear)output.textContent='';output.textContent+=text;output.scrollTop=output.scrollHeight;}
-function getPayload(){return{prompt:$('prompt').value,max_new_tokens:parseInt($('max_new_tokens').value),temperature:parseFloat($('temperature').value),steps:parseInt($('steps').value),block_size:parseInt($('block_size').value),cfg_scale:parseFloat($('cfg_scale').value),remasking:$('remasking').value,capture_interval:parseInt($('capture_interval').value)};}
-async function handleGenerate(){const payload=getPayload();const t0=performance.now();const res=await fetch('/generate',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload)});const data=await res.json();const ms=Math.round(performance.now()-t0);if(res.ok){appendOutput(`[${ms}ms]\n${data.generated_text||JSON.stringify(data,null,2)}\n\n`,true);setStatus(`OK — ${ms}ms`);}else{appendOutput(`Error ${res.status}:\n${JSON.stringify(data,null,2)}\n\n`,true);setStatus(`HTTP ${res.status}`,true);}}
-async function handleGenerateStream(){const payload=getPayload();const t0=performance.now();const res=await fetch('/generate_stream',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload)});const data=await res.json();const ms=Math.round(performance.now()-t0);if(res.ok){let text=`[${ms}ms]\nGenerated text:\n${data.generated_text}\n\n`;if(data.intermediate_states&&data.intermediate_states.length){text+=`Intermediate states (${data.intermediate_states.length}):\n`;data.intermediate_states.forEach((s,i)=>{text+=`  Step ${s.step}: ${s.text.substring(0,120).replace(/\n/g,' ')}...\n`;});}appendOutput(text+'\n',true);setStatus(`OK — ${ms}ms, ${data.intermediate_states?.length||0} intermediates`);}else{appendOutput(`Error ${res.status}:\n${JSON.stringify(data,null,2)}\n\n`,true);setStatus(`HTTP ${res.status}`,true);}}
-async function handleGenerateSSE(){const payload=getPayload();const t0=performance.now();appendOutput('',true);setStatus('Connecting SSE...');const res=await fetch('/generate_sse',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(payload)});if(!res.ok){const data=await res.json().catch(()=>({}));appendOutput(`Error ${res.status}:\n${JSON.stringify(data,null,2)}`,true);setStatus(`HTTP ${res.status}`,true);return;}const reader=res.body.getReader();const decoder=new TextDecoder();let buffer='';let finalText='';let eventCount=0;while(true){const {done,value}=await reader.read();if(done)break;buffer+=decoder.decode(value,{stream:true});const lines=buffer.split('\n');buffer=lines.pop();for(const line of lines){if(!line.startsWith('data: '))continue;const jsonStr=line.slice(6).trim();if(!jsonStr)continue;try{const event=JSON.parse(jsonStr);eventCount++;if(event.type==='final'){finalText=event.text;const ms=Math.round(performance.now()-t0);appendOutput(`[${ms}ms | ${eventCount} events]\n${finalText}\n`,true);setStatus(`Done — ${ms}ms, ${eventCount} events, ${event.total_steps||'?'} steps`);}else if(event.type==='intermediate'||event.type==='token'){appendOutput(`${event.text}`,true);setStatus(`Streaming... (${eventCount} events)`);}}catch(e){}}}if(!finalText&&eventCount===0){setStatus('Stream ended with no events',true);}}
-sendBtn.addEventListener('click',async()=>{sendBtn.disabled=true;setStatus('Sending...');try{const route=routeSelect.value;if(route==='/generate')await handleGenerate();else if(route==='/generate_stream')await handleGenerateStream();else if(route==='/generate_sse')await handleGenerateSSE();}catch(err){appendOutput(`Network/JS Error:\n${err.message}\n\n`,true);setStatus(err.message,true);}finally{sendBtn.disabled=false;}});
-</script>
-</body>
-</html>"""
 @app.route('/')
 def index():
     return Response(HTML_UI, mimetype='text/html')
 if __name__ == '__main__':
     load_model()
     app.run(host='0.0.0.0', port=int(os.getenv('PORT', 7860)))

 from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
+HTML_UI = """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>LLM API Tester</title>
+    <style>
+        * { box-sizing: border-box; }
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
+            max-width: 900px;
+            margin: 0 auto;
+            padding: 20px;
+            background: #f5f5f5;
+            color: #333;
+        }
+        h1 { margin-top: 0; font-size: 1.5rem; }
+        .card {
+            background: #fff;
+            border-radius: 8px;
+            padding: 20px;
+            margin-bottom: 16px;
+            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
+        }
+        label {
+            display: block;
+            font-weight: 600;
+            margin-bottom: 6px;
+            font-size: 0.9rem;
+        }
+        textarea, input, select {
+            width: 100%;
+            padding: 10px;
+            border: 1px solid #ddd;
+            border-radius: 6px;
+            font-size: 0.95rem;
+            font-family: inherit;
+        }
+        textarea { resize: vertical; min-height: 80px; }
+        .row {
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
+            gap: 12px;
+            margin-bottom: 12px;
+        }
+        .field { margin-bottom: 12px; }
+        .field.inline {
+            display: flex;
+            align-items: center;
+            gap: 8px;
+        }
+        .field.inline label { margin: 0; }
+        .field.inline input, .field.inline select {
+            width: auto;
+            flex: 1;
+        }
+        button {
+            background: #2563eb;
+            color: #fff;
+            border: none;
+            padding: 10px 20px;
+            border-radius: 6px;
+            font-size: 1rem;
+            cursor: pointer;
+            font-weight: 600;
+        }
+        button:hover { background: #1d4ed8; }
+        button:disabled { background: #93c5fd; cursor: not-allowed; }
+        .output {
+            background: #1e1e1e;
+            color: #e4e4e4;
+            padding: 16px;
+            border-radius: 6px;
+            font-family: ui-monospace, SFMono-Regular, "SF Mono", Menlo, Consolas, monospace;
+            font-size: 0.9rem;
+            white-space: pre-wrap;
+            word-break: break-word;
+            min-height: 120px;
+            max-height: 500px;
+            overflow-y: auto;
+        }
+        .output:empty::before {
+            content: "Response will appear here...";
+            color: #666;
+        }
+        .status {
+            font-size: 0.85rem;
+            color: #666;
+            margin-top: 8px;
+        }
+        .error { color: #dc2626; }
+        .success { color: #16a34a; }
+        .route-badge {
+            display: inline-block;
+            background: #e5e7eb;
+            padding: 2px 8px;
+            border-radius: 4px;
+            font-size: 0.8rem;
+            font-weight: 600;
+            margin-bottom: 8px;
+        }
+    </style>
+</head>
+<body>
+    <h1>LLM API Tester</h1>
+    <div class="card">
+        <div class="route-badge" id="routeBadge">/generate</div>
+        <div class="field">
+            <label for="route">Route</label>
+            <select id="route">
+                <option value="/generate">/generate (sync JSON)</option>
+                <option value="/generate_stream">/generate_stream (sync JSON + intermediates)</option>
+                <option value="/generate_sse">/generate_sse (streaming SSE)</option>
+            </select>
+        </div>
+        <div class="field">
+            <label for="prompt">Prompt</label>
+            <textarea id="prompt" placeholder="Enter your prompt here...">Write a short poem about stars</textarea>
+        </div>
+        <div class="row">
+            <div class="field">
+                <label for="max_new_tokens">max_new_tokens</label>
+                <input type="number" id="max_new_tokens" value="150" min="1" max="2048">
+            </div>
+            <div class="field">
+                <label for="temperature">temperature</label>
+                <input type="number" id="temperature" value="0.0" min="0" max="2" step="0.1">
+            </div>
+            <div class="field">
+                <label for="steps">steps (diffusion)</label>
+                <input type="number" id="steps" value="256" min="1">
+            </div>
+            <div class="field">
+                <label for="block_size">block_size</label>
+                <input type="number" id="block_size" value="32" min="1">
+            </div>
+        </div>
+        <div class="row">
+            <div class="field">
+                <label for="cfg_scale">cfg_scale</label>
+                <input type="number" id="cfg_scale" value="0.0" min="0" step="0.1">
+            </div>
+            <div class="field">
+                <label for="remasking">remasking</label>
+                <select id="remasking">
+                    <option value="low_confidence">low_confidence</option>
+                    <option value="random">random</option>
+                </select>
+            </div>
+            <div class="field">
+                <label for="capture_interval">capture_interval</label>
+                <input type="number" id="capture_interval" value="10" min="1">
+            </div>
+        </div>
+        <button id="sendBtn">Send Request</button>
+        <div class="status" id="status"></div>
+    </div>
+    <div class="card">
+        <label>Response</label>
+        <div class="output" id="output"></div>
+    </div>
+    <script>
+        const $ = id => document.getElementById(id);
+        const routeSelect = $('route');
+        const routeBadge = $('routeBadge');
+        const sendBtn = $('sendBtn');
+        const output = $('output');
+        const status = $('status');
+        routeSelect.addEventListener('change', () => {
+            routeBadge.textContent = routeSelect.value;
+        });
+        function setStatus(msg, isError = false) {
+            status.textContent = msg;
+            status.className = 'status ' + (isError ? 'error' : 'success');
+        }
+        function appendOutput(text, clear = false) {
+            if (clear) output.textContent = '';
+            output.textContent += text;
+            output.scrollTop = output.scrollHeight;
+        }
+        function getPayload() {
+            return {
+                prompt: $('prompt').value,
+                max_new_tokens: parseInt($('max_new_tokens').value),
+                temperature: parseFloat($('temperature').value),
+                steps: parseInt($('steps').value),
+                block_size: parseInt($('block_size').value),
+                cfg_scale: parseFloat($('cfg_scale').value),
+                remasking: $('remasking').value,
+                capture_interval: parseInt($('capture_interval').value)
+            };
+        }
+        async function handleGenerate() {
+            const payload = getPayload();
+            // Remove diffusion-only fields for non-diffusion if needed, but server ignores extras
+            const t0 = performance.now();
+            const res = await fetch('/generate', {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify(payload)
+            });
+            const data = await res.json();
+            const ms = Math.round(performance.now() - t0);
+            if (res.ok) {
+                appendOutput(`[${ms}ms]\n${data.generated_text || JSON.stringify(data, null, 2)}\n\n`, true);
+                setStatus(`OK — ${ms}ms`);
+            } else {
+                appendOutput(`Error ${res.status}:\n${JSON.stringify(data, null, 2)}\n\n`, true);
+                setStatus(`HTTP ${res.status}`, true);
+            }
+        }
+        async function handleGenerateStream() {
+            const payload = getPayload();
+            const t0 = performance.now();
+            const res = await fetch('/generate_stream', {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify(payload)
+            });
+            const data = await res.json();
+            const ms = Math.round(performance.now() - t0);
+            if (res.ok) {
+                let text = `[${ms}ms]\nGenerated text:\n${data.generated_text}\n\n`;
+                if (data.intermediate_states && data.intermediate_states.length) {
+                    text += `Intermediate states (${data.intermediate_states.length}):\n`;
+                    data.intermediate_states.forEach((s, i) => {
+                        text += "  Step " + s.step + ": " + s.text.substring(0,120).replace(/\\n/g, ' ') + "...\n";
+                    });
+                }
+                appendOutput(text + '\n', true);
+                setStatus(`OK — ${ms}ms, ${data.intermediate_states?.length || 0} intermediates`);
+            } else {
+                appendOutput(`Error ${res.status}:\n${JSON.stringify(data, null, 2)}\n\n`, true);
+                setStatus(`HTTP ${res.status}`, true);
+            }
+        }
+        async function handleGenerateSSE() {
+            const payload = getPayload();
+            const t0 = performance.now();
+            appendOutput('', true);
+            setStatus('Connecting SSE...');
+            const res = await fetch('/generate_sse', {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                body: JSON.stringify(payload)
+            });
+            if (!res.ok) {
+                const data = await res.json().catch(() => ({}));
+                appendOutput(`Error ${res.status}:\n${JSON.stringify(data, null, 2)}`, true);
+                setStatus(`HTTP ${res.status}`, true);
+                return;
+            }
+            const reader = res.body.getReader();
+            const decoder = new TextDecoder();
+            let buffer = '';
+            let finalText = '';
+            let eventCount = 0;
+            while (true) {
+                const { done, value } = await reader.read();
+                if (done) break;
+                buffer += decoder.decode(value, { stream: true });
+                const lines = buffer.split('\n');
+                buffer = lines.pop(); // keep incomplete line in buffer
+                for (const line of lines) {
+                    if (!line.startsWith('data: ')) continue;
+                    const jsonStr = line.slice(6).trim();
+                    if (!jsonStr) continue;
+                    try {
+                        const event = JSON.parse(jsonStr);
+                        eventCount++;
+                        if (event.type === 'final') {
+                            finalText = event.text;
+                            const ms = Math.round(performance.now() - t0);
+                            appendOutput(`[${ms}ms | ${eventCount} events]\n${finalText}\n`, true);
+                            setStatus(`Done — ${ms}ms, ${eventCount} events, ${event.total_steps || '?'} steps`);
+                        } else if (event.type === 'intermediate' || event.type === 'token') {
+                            // Live update: overwrite with latest accumulated text
+                            appendOutput(`${event.text}`, true);
+                            setStatus(`Streaming... (${eventCount} events)`);
+                        }
+                    } catch (e) {
+                        // ignore malformed lines
+                    }
+                }
+            }
+            if (!finalText && eventCount === 0) {
+                setStatus('Stream ended with no events', true);
+            }
+        }
+        sendBtn.addEventListener('click', async () => {
+            sendBtn.disabled = true;
+            setStatus('Sending...');
+            try {
+                const route = routeSelect.value;
+                if (route === '/generate') await handleGenerate();
+                else if (route === '/generate_stream') await handleGenerateStream();
+                else if (route === '/generate_sse') await handleGenerateSSE();
+            } catch (err) {
+                appendOutput(`Network/JS Error:\n${err.message}\n\n`, true);
+                setStatus(err.message, true);
+            } finally {
+                sendBtn.disabled = false;
+            }
+        });
+    </script>
+</body>
+</html>
+"""
 # 1. Environment Parsing & Architecture Strategy Mapping
 MODEL_NAME = os.getenv("MODEL_NAME", "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1")
 IS_DIFFUSION = "diffusion" in MODEL_NAME.lower()
     """Clone KV-cache. Fast path for tuples and Cache objects; falls back to deepcopy."""
     if pkv is None:
         return None
+    # Fast path: legacy tuple format
     if isinstance(pkv, tuple):
         return tuple(
             (k.clone() if k is not None else None, v.clone() if v is not None else None)
             for k, v in pkv
         )
+    # Fast path: transformers Cache objects (DynamicCache, etc.)
     if hasattr(pkv, 'key_cache') and hasattr(pkv, 'value_cache'):
         try:
             new_cache = pkv.__class__()
             return new_cache
         except Exception:
             pass
+    # Fallback
     return copy.deepcopy(pkv)
             torch_dtype=torch.bfloat16,
             trust_remote_code=False
         ).to(device).eval()
+    # Compile model for faster inference — ONLY for standard causal models
     if not IS_DIFFUSION:
         try:
             model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
         MODEL_NAME,
         trust_remote_code=IS_DIFFUSION
     )
+    print("Model compilation completed and loaded into memory workspace.")
 @app.route('/health', methods=['GET'])
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": prompt}
     ]
+    # enable_thinking=False for ALL routes to prevent Qwen3 from leaking internal monologue
     encoded = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
             accumulated = []
             for text in streamer:
+                if not text:  # skip empty chunks
                     continue
                 accumulated.append(text)
                 current = "".join(accumulated)
     )
 @app.route('/')
 def index():
     return Response(HTML_UI, mimetype='text/html')
 if __name__ == '__main__':
     load_model()
     app.run(host='0.0.0.0', port=int(os.getenv('PORT', 7860)))