diffusion-chatbot

Sleeping

App Files Files Community

Andhs commited on Jul 3

Commit

df717ea

verified ·

1 Parent(s): da77057

Upload app.py

Browse files

Files changed (1) hide show

app.py +332 -6

app.py CHANGED Viewed

@@ -9,7 +9,337 @@ import torch.nn.functional as F
 from flask import Flask, request, jsonify, Response
 from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModelForCausalLM, TextIteratorStreamer
 from threading import Thread
-from flask import redirect, send_from_directory
 # 1. Environment Parsing & Architecture Strategy Mapping
 MODEL_NAME = os.getenv("MODEL_NAME", "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1")
@@ -534,11 +864,7 @@ def generate_text_sse():
 @app.route('/')
 def index():
-    return redirect('/ui')
-@app.route('/ui')
-def ui():
-    return send_from_directory(os.path.dirname(os.path.abspath(__file__)), 'index.html')
 if __name__ == '__main__':
     load_model()

 from flask import Flask, request, jsonify, Response
 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")
 @app.route('/')
 def index():
+    return Response(HTML_UI, mimetype='text/html')
 if __name__ == '__main__':
     load_model()