Spaces:
Sleeping
Sleeping
| import sys | |
| import types | |
| import os | |
| import math | |
| import json | |
| import copy | |
| import torch | |
| import torch.nn.functional as F | |
| 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") | |
| IS_DIFFUSION = "diffusion" in MODEL_NAME.lower() | |
| # Dynamic initialization layer targeting Diffusion Language Models | |
| if IS_DIFFUSION: | |
| try: | |
| import dllm.utils | |
| import dllm.pipelines | |
| import dllm.data | |
| import dllm.core | |
| except ImportError: | |
| pass | |
| if 'dllm' not in sys.modules: | |
| dllm_mock = types.ModuleType('dllm') | |
| dllm_mock.core = sys.modules.get('dllm.core') | |
| dllm_mock.data = sys.modules.get('dllm.data') | |
| dllm_mock.pipelines = sys.modules.get('dllm.pipelines') | |
| dllm_mock.utils = sys.modules.get('dllm.utils') | |
| sys.modules['dllm'] = dllm_mock | |
| app = Flask(__name__) | |
| model = None | |
| tokenizer = None | |
| device = None | |
| # ========================================================== | |
| # SYSTEM WORKSPACE PIPELINES: CORE DIFFUSION SAMPLING LOOPS | |
| # ========================================================== | |
| def add_gumbel_noise(logits, temperature): | |
| """Add Gumbel noise using float32 (faster than float64 on most GPUs).""" | |
| if temperature == 0: | |
| return logits | |
| logits = logits.float() | |
| noise = torch.rand_like(logits) | |
| g = (-torch.log(noise)) ** temperature | |
| return logits.exp() / g | |
| def get_num_transfer_tokens(mask_index, steps): | |
| mask_num = mask_index.sum(dim=1, keepdim=True) | |
| base = mask_num // steps | |
| rem = mask_num % steps | |
| out = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.long) + base | |
| for i in range(mask_num.size(0)): | |
| out[i, : rem[i]] += 1 | |
| return out | |
| def build_staircase_attention_mask(x, block_size, pad_id): | |
| B, T = x.shape | |
| device = x.device | |
| valid = x != pad_id | |
| pos_raw = torch.cumsum(valid.long(), dim=-1) | |
| position_ids = torch.where(valid, pos_raw - 1, torch.zeros_like(pos_raw)).long() | |
| col = torch.arange(T, device=device) | |
| block_ids = (col // block_size).view(1, T).expand(B, T) | |
| block_ids = torch.where(valid, block_ids, torch.full_like(block_ids, -1)) | |
| q = block_ids.view(B, 1, T, 1) | |
| k = block_ids.view(B, 1, 1, T) | |
| attn = (k <= q) & (q >= 0) & (k >= 0) | |
| return attn, position_ids | |
| 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 | |
| # 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__() | |
| new_cache.key_cache = [k.clone() for k in pkv.key_cache] | |
| new_cache.value_cache = [v.clone() for v in pkv.value_cache] | |
| for attr in ('_seen_tokens', 'seen_tokens'): | |
| if hasattr(pkv, attr): | |
| setattr(new_cache, attr, getattr(pkv, attr)) | |
| return new_cache | |
| except Exception: | |
| pass | |
| # Fallback | |
| return copy.deepcopy(pkv) | |
| def diffusion_step_block(logits, x_block, mask_block, num_transfer, temperature, remasking): | |
| """Vectorized diffusion step — no per-sample Python loops.""" | |
| B, L, _ = logits.shape | |
| if not mask_block.any(): | |
| return x_block | |
| noisy = add_gumbel_noise(logits, temperature) | |
| x0 = noisy.argmax(dim=-1) | |
| if remasking == "low_confidence": | |
| p = F.softmax(logits, dim=-1) | |
| conf = p.gather(-1, x0.unsqueeze(-1)).squeeze(-1) | |
| elif remasking == "random": | |
| conf = torch.rand((B, L), device=logits.device) | |
| else: | |
| raise ValueError(remasking) | |
| x0 = torch.where(mask_block, x0, x_block) | |
| conf = conf.masked_fill(~mask_block, float("-inf")) | |
| k_max = int(num_transfer.max().item()) | |
| if k_max > 0: | |
| k = min(k_max, L) | |
| topk_vals, topk_idx = torch.topk(conf, k=k, dim=-1) | |
| commit = torch.zeros_like(x_block, dtype=torch.bool) | |
| valid_mask = torch.arange(k, device=x_block.device).view(1, k) < num_transfer.view(B, 1) | |
| commit.scatter_(1, topk_idx, valid_mask) | |
| x_block = torch.where(commit, x0, x_block) | |
| return x_block | |
| def generate(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size=32, | |
| temperature=0.0, cfg_scale=0.0, remasking="low_confidence", capture_interval=0): | |
| device = model.device | |
| mask_id = tokenizer.mask_token_id | |
| pad_id = tokenizer.pad_token_id | |
| if pad_id is None: | |
| pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.mask_token_id | |
| if isinstance(prompt, torch.Tensor): | |
| x = prompt.to(device).long() | |
| else: | |
| if isinstance(prompt[0], (list, tuple)): | |
| max_len = max(len(p) for p in prompt) | |
| x = torch.full((len(prompt), max_len), pad_id, device=device, dtype=torch.long) | |
| for i, p in enumerate(prompt): | |
| x[i, : len(p)] = torch.tensor(p, device=device) | |
| else: | |
| x = torch.tensor(prompt, device=device).long() | |
| if x.dim() == 1: | |
| x = x.unsqueeze(0) | |
| B = x.size(0) | |
| finished = torch.zeros(B, dtype=torch.bool, device=device) | |
| num_blocks = math.ceil(max_new_tokens / block_size) | |
| steps_per_block = math.ceil(steps / num_blocks) | |
| generated = 0 | |
| intermediates = [] | |
| total_step = 0 | |
| while generated < max_new_tokens: | |
| if finished.all(): | |
| break | |
| T_prefix = x.size(1) | |
| offset = T_prefix % block_size | |
| room = block_size if offset == 0 else block_size - offset | |
| cur_len = min(room, max_new_tokens - generated) | |
| if cur_len <= 0: | |
| break | |
| attn_pfx, pos_pfx = build_staircase_attention_mask(x, block_size, pad_id) | |
| out = model(x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True) | |
| cond_past = out.past_key_values | |
| if cfg_scale > 0: | |
| un_x = x.clone() | |
| un_x[:] = mask_id | |
| out_un = model(un_x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True) | |
| uncond_past = out_un.past_key_values | |
| else: | |
| uncond_past = None | |
| block = torch.full((B, cur_len), mask_id, device=device, dtype=torch.long) | |
| block[finished] = pad_id | |
| x = torch.cat([x, block], dim=1) | |
| T_total = x.size(1) | |
| block_mask = x[:, -cur_len:] == mask_id | |
| num_transfer = get_num_transfer_tokens(block_mask, steps_per_block) | |
| eff_steps = num_transfer.size(1) | |
| full_attn, full_pos = build_staircase_attention_mask(x, block_size, pad_id) | |
| attn_blk = full_attn[:, :, T_prefix:T_total, :] | |
| pos_blk = full_pos[:, T_prefix:T_total] | |
| for t in range(eff_steps): | |
| x_blk = x[:, T_prefix:T_total] | |
| m_blk = x_blk == mask_id | |
| cond_logits = model( | |
| x_blk, attention_mask=attn_blk, position_ids=pos_blk, | |
| past_key_values=clone_past_key_values(cond_past), use_cache=False | |
| ).logits | |
| logits = cond_logits | |
| if cfg_scale > 0: | |
| un_logits = model( | |
| x_blk, attention_mask=attn_blk, position_ids=pos_blk, | |
| past_key_values=clone_past_key_values(uncond_past), use_cache=False | |
| ).logits | |
| logits = un_logits + (cfg_scale + 1.0) * (cond_logits - un_logits) | |
| x_blk_new = diffusion_step_block( | |
| logits, x_blk, m_blk, num_transfer[:, t], temperature, remasking | |
| ) | |
| x[:, T_prefix:T_total] = x_blk_new | |
| if capture_interval > 0 and total_step % capture_interval == 0: | |
| intermediates.append(x.clone()) | |
| total_step += 1 | |
| if tokenizer.eos_token_id is not None: | |
| finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1) | |
| generated += cur_len | |
| if finished.all(): | |
| break | |
| if capture_interval > 0: | |
| return x, intermediates | |
| return x | |
| def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, block_size=32, | |
| temperature=0.0, cfg_scale=0.0, remasking="low_confidence", capture_interval=10): | |
| device = model.device | |
| mask_id = tokenizer.mask_token_id | |
| pad_id = tokenizer.pad_token_id | |
| if pad_id is None: | |
| pad_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else tokenizer.mask_token_id | |
| if isinstance(prompt, torch.Tensor): | |
| x = prompt.to(device).long() | |
| else: | |
| if isinstance(prompt[0], (list, tuple)): | |
| max_len = max(len(p) for p in prompt) | |
| x = torch.full((len(prompt), max_len), pad_id, device=device, dtype=torch.long) | |
| for i, p in enumerate(prompt): | |
| x[i, : len(p)] = torch.tensor(p, device=device) | |
| else: | |
| x = torch.tensor(prompt, device=device).long() | |
| if x.dim() == 1: | |
| x = x.unsqueeze(0) | |
| B = x.size(0) | |
| finished = torch.zeros(B, dtype=torch.bool, device=device) | |
| num_blocks = math.ceil(max_new_tokens / block_size) | |
| steps_per_block = math.ceil(steps / num_blocks) | |
| generated = 0 | |
| total_step = 0 | |
| prompt_len = x.size(1) | |
| while generated < max_new_tokens: | |
| if finished.all(): | |
| break | |
| T_prefix = x.size(1) | |
| offset = T_prefix % block_size | |
| room = block_size if offset == 0 else block_size - offset | |
| cur_len = min(room, max_new_tokens - generated) | |
| if cur_len <= 0: | |
| break | |
| attn_pfx, pos_pfx = build_staircase_attention_mask(x, block_size, pad_id) | |
| out = model(x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True) | |
| cond_past = out.past_key_values | |
| if cfg_scale > 0: | |
| un_x = x.clone() | |
| un_x[:] = mask_id | |
| out_un = model(un_x, attention_mask=attn_pfx, position_ids=pos_pfx, use_cache=True) | |
| uncond_past = out_un.past_key_values | |
| else: | |
| uncond_past = None | |
| block = torch.full((B, cur_len), mask_id, device=device, dtype=torch.long) | |
| block[finished] = pad_id | |
| x = torch.cat([x, block], dim=1) | |
| T_total = x.size(1) | |
| block_mask = x[:, -cur_len:] == mask_id | |
| num_transfer = get_num_transfer_tokens(block_mask, steps_per_block) | |
| eff_steps = num_transfer.size(1) | |
| full_attn, full_pos = build_staircase_attention_mask(x, block_size, pad_id) | |
| attn_blk = full_attn[:, :, T_prefix:T_total, :] | |
| pos_blk = full_pos[:, T_prefix:T_total] | |
| for t in range(eff_steps): | |
| x_blk = x[:, T_prefix:T_total] | |
| m_blk = x_blk == mask_id | |
| cond_logits = model( | |
| x_blk, attention_mask=attn_blk, position_ids=pos_blk, | |
| past_key_values=clone_past_key_values(cond_past), use_cache=False | |
| ).logits | |
| logits = cond_logits | |
| if cfg_scale > 0: | |
| un_logits = model( | |
| x_blk, attention_mask=attn_blk, position_ids=pos_blk, | |
| past_key_values=clone_past_key_values(uncond_past), use_cache=False | |
| ).logits | |
| logits = un_logits + (cfg_scale + 1.0) * (cond_logits - un_logits) | |
| x_blk_new = diffusion_step_block( | |
| logits, x_blk, m_blk, num_transfer[:, t], temperature, remasking | |
| ) | |
| x[:, T_prefix:T_total] = x_blk_new | |
| if total_step % capture_interval == 0: | |
| new_tokens = x[0, prompt_len:prompt_len + max_new_tokens].tolist() | |
| text = tokenizer.decode(new_tokens, skip_special_tokens=True) | |
| yield { | |
| "type": "intermediate", | |
| "step": total_step, | |
| "text": text, | |
| "total_steps": steps | |
| } | |
| total_step += 1 | |
| if tokenizer.eos_token_id is not None: | |
| finished |= (x_blk_new == tokenizer.eos_token_id).any(dim=1) | |
| if finished.all(): | |
| break | |
| generated += cur_len | |
| if finished.all(): | |
| break | |
| new_tokens = x[0, prompt_len:prompt_len + max_new_tokens].tolist() | |
| final_text = tokenizer.decode(new_tokens, skip_special_tokens=True) | |
| yield { | |
| "type": "final", | |
| "text": final_text, | |
| "total_steps": total_step | |
| } | |
| # ========================================================== | |
| # ARCHITECTURE ROUTING LAYERS & TRANSLATION ENGINE CODES | |
| # ========================================================== | |
| def load_model(): | |
| global model, tokenizer, device | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Initializing {MODEL_NAME} on {device}... (Diffusion Strategy Flag = {IS_DIFFUSION})") | |
| if IS_DIFFUSION: | |
| model = AutoModelForMaskedLM.from_pretrained( | |
| MODEL_NAME, | |
| torch_dtype=torch.bfloat16, | |
| trust_remote_code=True | |
| ).to(device).eval() | |
| else: | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_NAME, | |
| 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) | |
| print("Model compiled with torch.compile.") | |
| except Exception as e: | |
| print(f"torch.compile skipped: {e}") | |
| else: | |
| print("Diffusion model loaded without torch.compile (custom FX code incompatible with Dynamo).") | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| MODEL_NAME, | |
| trust_remote_code=IS_DIFFUSION | |
| ) | |
| print("Model compilation completed and loaded into memory workspace.") | |
| def health(): | |
| return jsonify({"status": "healthy", "model_loaded": model is not None, "is_diffusion": IS_DIFFUSION}) | |
| def generate_text(): | |
| if model is None or tokenizer is None: | |
| return jsonify({"error": "Model initialization missing"}), 503 | |
| data = request.get_json() or {} | |
| if 'prompt' not in data: | |
| return jsonify({"error": "Missing 'prompt' operational field"}), 400 | |
| prompt = data['prompt'] | |
| max_new_tokens = data.get('max_new_tokens', 256) | |
| temperature = data.get('temperature', 0.0) | |
| system_prompt = data.get('system_prompt', 'You are an expert real-time translation assistant.') | |
| messages = [ | |
| {"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, | |
| tokenize=True, | |
| enable_thinking=False | |
| ) | |
| if IS_DIFFUSION: | |
| input_ids = torch.tensor([encoded], dtype=torch.long, device=device) | |
| steps = data.get('steps', 256) | |
| block_size = data.get('block_size', 32) | |
| cfg_scale = data.get('cfg_scale', 0.0) | |
| remasking = data.get('remasking', 'low_confidence') | |
| output = generate( | |
| model, tokenizer, input_ids, | |
| steps=steps, max_new_tokens=max_new_tokens, block_size=block_size, | |
| temperature=temperature, cfg_scale=cfg_scale, remasking=remasking, | |
| ) | |
| prompt_len = len(encoded) | |
| new_tokens = output[0, prompt_len:prompt_len + max_new_tokens].tolist() | |
| generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True) | |
| else: | |
| input_ids = torch.tensor([encoded], dtype=torch.long, device=device) | |
| output_ids = model.generate( | |
| input_ids, | |
| max_new_tokens=max_new_tokens, | |
| temperature=temperature, | |
| do_sample=True if temperature > 0 else False, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| generated_ids = output_ids[0, input_ids.shape[-1]:] | |
| generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True) | |
| return jsonify({"prompt": prompt, "generated_text": generated_text}) | |
| def generate_text_stream(): | |
| if model is None or tokenizer is None: | |
| return jsonify({"error": "Model workspace offline"}), 503 | |
| data = request.get_json() or {} | |
| if not data or 'prompt' not in data: | |
| return jsonify({"error": "Missing 'prompt' operational field"}), 400 | |
| prompt = data['prompt'] | |
| max_new_tokens = data.get('max_new_tokens', 256) | |
| temperature = data.get('temperature', 0.0) | |
| system_prompt = data.get('system_prompt', 'You are an expert real-time translation assistant.') | |
| messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": prompt} | |
| ] | |
| encoded = tokenizer.apply_chat_template( | |
| messages, add_generation_prompt=True, tokenize=True, | |
| enable_thinking=False | |
| ) | |
| if IS_DIFFUSION: | |
| input_ids = torch.tensor([encoded], dtype=torch.long, device=device) | |
| steps = data.get('steps', 256) | |
| block_size = data.get('block_size', 32) | |
| cfg_scale = data.get('cfg_scale', 0.0) | |
| remasking = data.get('remasking', 'low_confidence') | |
| capture_interval = data.get('capture_interval', 10) | |
| output, intermediates = generate( | |
| model, tokenizer, input_ids, | |
| steps=steps, max_new_tokens=max_new_tokens, block_size=block_size, | |
| temperature=temperature, cfg_scale=cfg_scale, remasking=remasking, | |
| capture_interval=capture_interval, | |
| ) | |
| prompt_len = len(encoded) | |
| intermediate_states = [] | |
| for i, intermediate in enumerate(intermediates): | |
| new_tokens = intermediate[0, prompt_len:prompt_len + max_new_tokens].tolist() | |
| text = tokenizer.decode(new_tokens, skip_special_tokens=True) | |
| intermediate_states.append({"step": i * capture_interval, "text": text}) | |
| new_tokens = output[0, prompt_len:prompt_len + max_new_tokens].tolist() | |
| generated_text = tokenizer.decode(new_tokens, skip_special_tokens=True) | |
| return jsonify({"prompt": prompt, "generated_text": generated_text, "intermediate_states": intermediate_states}) | |
| else: | |
| input_ids = torch.tensor([encoded], dtype=torch.long, device=device) | |
| output_ids = model.generate( | |
| input_ids, max_new_tokens=max_new_tokens, temperature=temperature, | |
| do_sample=True if temperature > 0 else False, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| generated_ids = output_ids[0, input_ids.shape[-1]:] | |
| generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True) | |
| return jsonify({"prompt": prompt, "generated_text": generated_text, "intermediate_states": []}) | |
| def generate_text_sse(): | |
| if model is None or tokenizer is None: | |
| return jsonify({"error": "Model workspace offline"}), 503 | |
| data = request.get_json() or {} | |
| if not data or 'prompt' not in data: | |
| return jsonify({"error": "Missing 'prompt' operational field"}), 400 | |
| prompt = data['prompt'] | |
| max_new_tokens = data.get('max_new_tokens', 256) | |
| temperature = data.get('temperature', 0.0) | |
| system_prompt = data.get('system_prompt', 'You are an expert real-time translation assistant.') | |
| messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| {"role": "user", "content": prompt} | |
| ] | |
| encoded = tokenizer.apply_chat_template( | |
| messages, add_generation_prompt=True, tokenize=True, | |
| enable_thinking=False | |
| ) | |
| input_ids = torch.tensor([encoded], dtype=torch.long, device=device) | |
| def stream(): | |
| if IS_DIFFUSION: | |
| steps = data.get('steps', 256) | |
| block_size = data.get('block_size', 32) | |
| cfg_scale = data.get('cfg_scale', 0.0) | |
| remasking = data.get('remasking', 'low_confidence') | |
| capture_interval = data.get('capture_interval', 10) | |
| for state in generate_stream( | |
| model, tokenizer, input_ids, | |
| steps=steps, max_new_tokens=max_new_tokens, block_size=block_size, | |
| temperature=temperature, cfg_scale=cfg_scale, remasking=remasking, | |
| capture_interval=capture_interval | |
| ): | |
| yield f"data: {json.dumps(state)}\n\n" | |
| else: | |
| streamer = TextIteratorStreamer( | |
| tokenizer, skip_prompt=True, skip_special_tokens=True | |
| ) | |
| generation_kwargs = dict( | |
| input_ids=input_ids, | |
| streamer=streamer, | |
| max_new_tokens=max_new_tokens, | |
| temperature=temperature, | |
| do_sample=True if temperature > 0 else False, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| def _generate(): | |
| with torch.inference_mode(): | |
| model.generate(**generation_kwargs) | |
| thread = Thread(target=_generate) | |
| thread.start() | |
| accumulated = [] | |
| for text in streamer: | |
| if not text: # skip empty chunks | |
| continue | |
| accumulated.append(text) | |
| current = "".join(accumulated) | |
| yield f"data: {json.dumps({'type': 'intermediate', 'text': current})}\n\n" | |
| full_text = "".join(accumulated) | |
| yield f"data: {json.dumps({'type': 'final', 'text': full_text, 'total_steps': 1})}\n\n" | |
| return Response( | |
| stream(), mimetype='text/event-stream', | |
| headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'} | |
| ) | |
| 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))) | |