import gradio as gr import torch import numpy as np import json import time from html import escape as html_escape from transformers import AutoTokenizer import os import importlib import os from huggingface_hub import hf_hub_download import spaces from dotenv import load_dotenv from infer import ( load_trained_model, find_answer_start, get_noising_schedule, noisify_answer, filter_logits, confidence_guided_noising, noisify_answer_without_remasking ) from models import CustomTransformerModel from model_config import CustomTransformerConfig # Load .env only when running locally if os.getenv("HF_TOKEN") is None: load_dotenv() hf_token = os.getenv("HF_TOKEN") if hf_token is None: raise ValueError("HF_TOKEN is not set") rng = np.random.default_rng() def _generate_step(input_ids, top_p, top_k, eos_bias=0.0): """Single diffusion inference step. Must be called inside a @spaces.GPU context.""" with torch.no_grad(): input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device) with torch.autocast(device_type=model.device.type, dtype=torch.float16): logits = model(input_ids=input_tensor)["logits"] # Apply eos_bias if eos_bias != 0.0: logits[0, :, eos_token_id] += eos_bias logits = filter_logits(logits, top_k=top_k, top_p=top_p) logits = logits.clamp(min=-1e8, max=1e4) probs = torch.nn.functional.softmax(logits, dim=-1)[0] probs = torch.clamp(probs, min=1e-8, max=1.0) assert torch.all(torch.isfinite(probs)), "Non-finite values in probs!" assert (probs >= 0).all(), "Negative probs!" sampled = torch.multinomial(probs, num_samples=1).squeeze(-1).tolist() # Extract confidence of selected tokens conf = probs[range(len(sampled)), sampled].cpu().numpy() return sampled, conf @spaces.GPU def run_diffusion_loop(input_ids, answer_start, max_it, pause_length, eos_bias, sharpness, noise_start, use_confidence_noising, use_permanent_unmasking, noise_clipping, top_p, top_k, added_tokens, context_window): """Full diffusion loop held inside a single GPU context to avoid repeated re-acquisition.""" ori_input_tokens = input_ids[:] current_tokens, just_noised_indices = noisify_answer( input_ids, answer_start, tokenizer, threshold=1.0, noise_start=1.0 ) yield render_html("Iteration 0 (initial noise)", highlight_tokens(current_tokens[answer_start:], answer_start, just_noised_indices, color="red")), None if pause_length > 0: time.sleep(pause_length) last_tokens = [] prev_decoded = [] unmasked_mask = [False] * len(current_tokens) current_tokens = current_tokens[:answer_start] for i in range(max_it): current_tokens = current_tokens + [mask_token_id] * added_tokens current_tokens = current_tokens[:context_window] generated_tokens, confidences = _generate_step(current_tokens, top_p, top_k, eos_bias=eos_bias) current_tokens = ori_input_tokens[:answer_start] + generated_tokens[answer_start:] new_decoded = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:]) diff_indices = { answer_start + j for j, tok in enumerate(new_decoded) if j >= len(prev_decoded) or tok != prev_decoded[j] } prev_decoded = new_decoded yield render_html(f"Iteration {i+1}/{max_it} (after generation)", highlight_tokens(current_tokens[answer_start:], answer_start, diff_indices, color="green")), None if pause_length > 0: time.sleep(pause_length) last_tokens.append(current_tokens) if len(last_tokens) > 3: last_tokens.pop(0) if len(last_tokens) == 3 and last_tokens[0] == last_tokens[1] == last_tokens[2]: yield render_html("Stopped early", f"After {i+1} iterations."), None break if i < max_it - 1: threshold = get_noising_schedule(i, max_it, sharpness=sharpness) if use_confidence_noising: noised_answer, just_noised_indices = confidence_guided_noising( current_tokens, answer_start, tokenizer, confidences, noise_clipping, threshold=threshold, noise_start=noise_start, unmasked_mask=unmasked_mask if use_permanent_unmasking else None ) elif use_permanent_unmasking: noised_answer, just_noised_indices = noisify_answer_without_remasking( current_tokens, answer_start, tokenizer, threshold=threshold, noise_start=noise_start, unmasked_mask=unmasked_mask ) else: noised_answer, just_noised_indices = noisify_answer( current_tokens, answer_start, tokenizer, threshold=threshold, noise_start=noise_start ) for idx in range(answer_start, len(current_tokens)): if noised_answer[idx] != mask_token_id: unmasked_mask[idx] = True if use_permanent_unmasking: permanently_unmasked = {idx for idx in range(answer_start, len(current_tokens)) if unmasked_mask[idx]} yield render_html(f"Iteration {i+1}/{max_it} (after noising)", highlight_tokens(noised_answer[answer_start:], answer_start, permanently_unmasked, color="green")), None else: yield render_html(f"Iteration {i+1}/{max_it} (after noising)", highlight_tokens(noised_answer[answer_start:], answer_start, just_noised_indices, color="red")), None if pause_length > 0: time.sleep(pause_length) current_tokens = ori_input_tokens[:answer_start] + noised_answer[answer_start:] answer_ids = current_tokens[answer_start:] # Strip trailing EOS and MASK tokens before decoding while answer_ids and answer_ids[-1] in (eos_token_id, mask_token_id): answer_ids = answer_ids[:-1] try: final_ids = answer_ids[:answer_ids.index(eos_token_id)] except ValueError: final_ids = answer_ids final_output = tokenizer.decode(final_ids, skip_special_tokens=True).lstrip() yield render_html(f"Final Output ({len(final_ids)} tokens after {i+1} iterations)", final_output), final_output # type: ignore def format_chat_prompt(question): return ( "<|begin_of_text|>\n" "<|start_header_id|>system<|end_header_id|>\n" "You are a helpful assistant.\n" "<|start_header_id|>user<|end_header_id|>\n" f"{question}\n" "<|start_header_id|>assistant<|end_header_id|>\n" ) def format_multiturn_prompt(history, question): parts = [ "<|begin_of_text|>\n" "<|start_header_id|>system<|end_header_id|>\n" "You are a helpful assistant.\n" ] for q, a in history: parts.append(f"<|start_header_id|>user<|end_header_id|>\n{q}\n") parts.append(f"<|start_header_id|>assistant<|end_header_id|>\n{a}\n") parts.append(f"<|start_header_id|>user<|end_header_id|>\n{question}\n") parts.append("<|start_header_id|>assistant<|end_header_id|>\n") return "".join(parts) def render_html(label, text): return f"{label}
{text.lstrip()}
" def render_chat_history(history): if not history: return "" parts = [] for q, a in history: parts.append( f'
' f'
' f'{html_escape(q.strip())}
' ) parts.append( f'
' f'
' f'{html_escape(a.strip())}
' ) content = "".join(parts) return ( f'
' f'{content}
' f'' ) def highlight_tokens(token_ids, answer_start, changed_indices, color): # Strip trailing EOS and MASK tokens so the tail is invisible end = len(token_ids) while end > 0 and token_ids[end - 1] in (eos_token_id, mask_token_id): end -= 1 token_ids = token_ids[:end] tokens = tokenizer.convert_ids_to_tokens(token_ids) highlighted = [] for j, tok in enumerate(tokens): tok_id = tokenizer.convert_tokens_to_ids(tok) if tok_id == eos_token_id: continue if tok_id == mask_token_id: highlighted.append('mask') continue tok_str = tokenizer.convert_tokens_to_string([tok]) if (answer_start + j) in changed_indices: highlighted.append(f'{tok_str}') else: highlighted.append(tok_str) return "".join(highlighted) def on_generate(question, history, multiturn_on, max_it, pause_length, eos_bias, sharpness, noise_start, use_confidence_noising, use_permanent_unmasking, noise_clipping, top_p, top_k, added_tokens, context_window): eos_bias = -eos_bias q = question.strip() or "What do you know about the city of Amsterdam?" prompt = format_multiturn_prompt(history, q) if (multiturn_on and history) else format_chat_prompt(q) input_ids = tokenizer.encode(prompt, add_special_tokens=False) answer_start = find_answer_start(input_ids, assistant_marker_ids) if answer_start is None: yield history, render_html("Error", "Could not find Assistant marker in input.") return input_ids = (input_ids + [mask_token_id] * (context_window - len(input_ids)))[:context_window] history_html = render_chat_history(history) if multiturn_on and history else "" final_answer = None last_step_html = "" for step_html, final in run_diffusion_loop( input_ids, answer_start, max_it, pause_length, eos_bias, sharpness, noise_start, use_confidence_noising, use_permanent_unmasking, noise_clipping, top_p, top_k, added_tokens, context_window ): if final is not None: final_answer = final last_step_html = step_html yield history, history_html + step_html # Generation complete — update history and switch to chat view if multi-turn if multiturn_on and final_answer is not None: new_history = history + [(q, final_answer)] yield new_history, render_chat_history(new_history) else: yield history, last_step_html def is_running_on_spaces(): return os.getenv("SPACE_ID") is not None print("Loading model...") if is_running_on_spaces(): # Load from Hugging Face Hub ckpt_path = hf_hub_download( repo_id="ruurd/tini_model", filename="diffusion-model-8B.pth", token=os.getenv("HF_TOKEN") ) else: # Load from local path ckpt_path = "diffusion-model-8B.pth" # change to your actual local path model, tokenizer = load_trained_model(checkpoint_path=ckpt_path) print("✅ Model loaded.") vocab_size = len(tokenizer) eos_token_id = tokenizer.eos_token_id mask_token_id = tokenizer.encode('MASK', add_special_tokens=False)[0] assistant_marker_ids = tokenizer.encode("<|start_header_id|>assistant<|end_header_id|>", add_special_tokens=False) with gr.Blocks(title="Diffusion Language Model Chat") as demo: chat_history_state = gr.State([]) gr.Markdown("## Diffusion Language Model Chat") gr.Markdown("A diffusion-based language model that generates answers progressively.") with gr.Row(): with gr.Column(scale=1): question_input = gr.Textbox( label="User Question", lines=2, placeholder="What do you know about the city of Amsterdam?" ) with gr.Row(): submit_btn = gr.Button("Generate", variant="primary") stop_btn = gr.Button("Stop") clear_btn = gr.Button("Clear Chat") multiturn_checkbox = gr.Checkbox(label="Multi-turn dialog", value=False) max_it_slider = gr.Slider(1, 512, value=64, step=1, label="Number of iterations: ↑ = more iterations") pause_slider = gr.Slider(0, 5, value=0.01, step=0.01, label="Pause between iterations ↑ = longer pause") eos_bias_slider = gr.Slider(-5.0, 5.0, value=0.0, step=0.1, label="Generation length: ↑ = more output tokens by decreasing eos token probability") sharpness_slider = gr.Slider(1.0, 20.0, value=1.0, step=0.5, label="Noise decay sharpness: ↓ = more noise in later iterations") noise_start_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Noise start fraction: ↑ = more noise") confidence_noising_cb = gr.Checkbox(value=False, label="Use confidence-guided noising") permanent_unmasking_cb= gr.Checkbox(value=False, label="Use permanent unmasking") noise_clipping_slider = gr.Slider(0.01, 1.0, value=0.01, step=0.01, label="Noise clipping: ↓ = more confidence guidance") topk_slider = gr.Slider(1, 1000, value=3, step=1, label="Top-k: ↑ = more random answers") topp_slider = gr.Slider(0.0, 1.0, value=1.0, step=0.01, label="Top-p: ↑ = more random answers") added_tokens_slider = gr.Slider(1, 256, value=256, step=1, label="Semi-autoregressive generation: number of added tokens per iteration") context_window_slider = gr.Slider(128, 2048, value=256, step=1, label="Context window: ↑ = longer sequences") with gr.Column(scale=1): output_html = gr.HTML(label="Output") all_inputs = [ question_input, chat_history_state, multiturn_checkbox, max_it_slider, pause_slider, eos_bias_slider, sharpness_slider, noise_start_slider, confidence_noising_cb, permanent_unmasking_cb, noise_clipping_slider, topp_slider, topk_slider, added_tokens_slider, context_window_slider, ] all_outputs = [chat_history_state, output_html] gen_event = submit_btn.click(on_generate, inputs=all_inputs, outputs=all_outputs) sub_event = question_input.submit(on_generate, inputs=all_inputs, outputs=all_outputs) stop_btn.click(fn=None, inputs=None, outputs=None, cancels=[gen_event, sub_event]) def on_clear(): return [], "" clear_btn.click(on_clear, inputs=[], outputs=[chat_history_state, output_html]) def on_toggle(multiturn_on, history): return render_chat_history(history) if multiturn_on else "" multiturn_checkbox.change( on_toggle, inputs=[multiturn_checkbox, chat_history_state], outputs=[output_html] ) demo.launch(share=True, allowed_paths=["."], ssr_mode=False)