import gradio as gr import torch import numpy as np import json import time from transformers import AutoTokenizer from llama_diffusion_model import disable_dropout import os import importlib from huggingface_hub import hf_hub_download from llama_diffusion_model import CustomTransformerModel, CustomTransformerConfig, BidirectionalLlamaAttention, disable_dropout import spaces hf_token = os.getenv("HF_TOKEN") # --- Load tokenizer --- tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B", use_fast=True, token=hf_token) vocab_size = len(tokenizer) pad_token = tokenizer.pad_token_id or tokenizer.eos_token_id eot_token_id = tokenizer.eos_token_id assistant_marker_ids = tokenizer.encode("Assistant:", add_special_tokens=False) # --- Load token probabilities --- with open("token_probabilities.json") as f: token_probs_dict = json.load(f) token_probabilities = np.array([token_probs_dict[str(i)] for i in range(len(token_probs_dict))], dtype=np.float32) def load_model(): ckpt_path = hf_hub_download( repo_id="ruurd/tini_model", filename="diffusion-model.pth", token=os.getenv("HF_TOKEN") ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = torch.load(ckpt_path, map_location=device) model = disable_dropout(model) model.to(device) model.eval() return model rng = np.random.default_rng() # --- Utility Functions --- def decode_tokens_safe(token_ids): return tokenizer.decode(token_ids, skip_special_tokens=True).replace("\n", " ") def find_answer_start(input_ids, marker_ids): for i in range(len(input_ids) - len(marker_ids) + 1): if input_ids[i:i + len(marker_ids)] == marker_ids: return i + len(marker_ids) return None def get_noising_schedule(i, max_it, sharpness=5.0): x = i / max_it return (np.exp(-sharpness * x) - np.exp(-sharpness)) / (1 - np.exp(-sharpness)) def noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=1.0): noised = input_ids.copy() answer_len = len(input_ids) - answer_start num_to_noise = int(threshold * answer_len) if num_to_noise > 0: indices = rng.choice(np.arange(answer_start, len(input_ids)), size=num_to_noise, replace=False) mixed_probs = token_probabilities.copy() mixed_probs[eot_token_id] *= eot_weight mixed_probs /= mixed_probs.sum() noise = rng.choice(np.arange(vocab_size), size=num_to_noise, p=mixed_probs) for idx, val in zip(indices, noise): noised[idx] = val return noised def generate_diffusion_text(input_ids, answer_start): with torch.no_grad(): input_tensor = torch.tensor([input_ids], dtype=torch.long).to(model.device) logits = model(input_ids=input_tensor)["logits"] probs = torch.nn.functional.softmax(logits, dim=-1).squeeze() probs = torch.clamp(probs, min=1e-8, max=1.0) sampled = torch.multinomial(probs, num_samples=1).squeeze().tolist() return input_ids[:answer_start] + sampled[answer_start:] # --- Inference Wrapper --- # --- Gradio Interface --- print("Loading model...") model = load_model() print("✅ Model loaded.") # --- Generation logic --- @spaces.GPU def run_diffusion_loop(question, eot_weight, max_it, sharpness): placeholder = "What do you know about the city of New York?" if question.strip() == "": question = placeholder prompt = f"User: {question}\nAssistant:" input_ids = tokenizer.encode(prompt, add_special_tokens=False) answer_start = find_answer_start(input_ids, assistant_marker_ids) if answer_start is None: return [], "Error: Could not find Assistant marker in input." input_ids = (input_ids + [pad_token] * (256 - len(input_ids)))[:256] current_tokens = noisify_answer(input_ids, answer_start, threshold=1.0, eot_weight=eot_weight) prev_decoded_tokens = [] last_tokens = [] history = ["**User:** " + question] for i in range(max_it): generated_tokens = generate_diffusion_text(current_tokens, answer_start) current_tokens = generated_tokens decoded_ids = current_tokens[answer_start:] decoded_tokens = tokenizer.convert_ids_to_tokens(decoded_ids) filtered_tokens = [tok for tok in decoded_tokens if tokenizer.convert_tokens_to_ids(tok) != eot_token_id] filtered_prev_tokens = [tok for tok in prev_decoded_tokens if tokenizer.convert_tokens_to_ids(tok) != eot_token_id] if prev_decoded_tokens else [] highlighted = [] for tok_new, tok_old in zip(filtered_tokens, filtered_prev_tokens): text = tokenizer.convert_tokens_to_string([tok_new]) if tok_new != tok_old: highlighted.append(f"{text}") else: highlighted.append(text) prev_decoded_tokens = decoded_tokens last_tokens.append(generated_tokens) if len(last_tokens) == 3 and all(t == last_tokens[0] for t in last_tokens): break threshold = get_noising_schedule(i, max_it, sharpness=sharpness) current_tokens = noisify_answer(generated_tokens, answer_start, threshold=threshold, eot_weight=eot_weight) time.sleep(0.01) final_tokens = tokenizer.convert_ids_to_tokens(current_tokens[answer_start:]) final_tokens = [tok for tok in final_tokens if tokenizer.convert_tokens_to_ids(tok) != eot_token_id] final_output = tokenizer.convert_tokens_to_string(final_tokens) history.append("**Assistant:** " + final_output) return history, final_output # --- UI Layout --- css = ".category-legend{display:none}" with gr.Blocks(css=css) as demo: gr.Markdown("# Tini Diffusion LLM 🌀") with gr.Row(): with gr.Column(scale=3): chatbox = gr.Chatbot(label="Conversation", value=[], height=400) question_input = gr.Textbox(label="Your Question", placeholder="What do you want to ask?", scale=8) send_btn = gr.Button("Generate") with gr.Column(scale=2): eot_weight = gr.Slider(0, 1, value=0.4, step=0.05, label="EOT weight") max_iters = gr.Slider(1, 512, value=64, step=1, label="Iterations") sharpness = gr.Slider(1.0, 20.0, value=5.0, step=0.5, label="Sharpness") def handle_submit(question, eot, max_it, sharp): history, _ = run_diffusion_loop(question, eot, max_it, sharp) return history send_btn.click( fn=handle_submit, inputs=[question_input, eot_weight, max_iters, sharpness], outputs=[chatbox] ) demo.queue().launch()