Spaces:
Running
Running
File size: 2,303 Bytes
f301955 f6b3246 4b51c4b f301955 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import torch
import gradio as gr
from modeling_diffusion import DiffusionTextModel
# =====================
# Load Model from Hub
# =====================
model = DiffusionTextModel.from_pretrained("yasserrmd/diffusion-text-demo")
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
PAD_TOKEN = "[PAD]"
MASK_TOKEN = "[MASK]"
vocab = {PAD_TOKEN: 0, MASK_TOKEN: 1}
# Reverse mapping
id_to_word = {i: w for w, i in vocab.items()}
# Special token IDs
pad_id = vocab[PAD_TOKEN]
mask_id = vocab[MASK_TOKEN]
# =====================
# Generation Function
# =====================
def generate_with_prompt(model, input_text, max_length=50, T=10):
# Ensure max_length does not exceed 99
max_length = min(max_length, 99)
model.eval()
input_tokens = input_text.split()
input_ids = [vocab.get(tok, mask_id) for tok in input_tokens]
seq = torch.full((1, max_length), mask_id, dtype=torch.long, device=device)
seq[0, :len(input_ids)] = torch.tensor(input_ids, device=device)
for step in range(T, 0, -1):
with torch.no_grad():
logits = model(seq, torch.tensor([step], device=device))
probs = torch.softmax(logits, dim=-1)
for pos in range(len(input_ids), max_length):
if seq[0, pos].item() == mask_id:
seq[0, pos] = torch.multinomial(probs[0, pos], 1)
ids = seq[0].tolist()
if pad_id in ids:
ids = ids[:ids.index(pad_id)]
return " ".join(id_to_word[i] for i in ids)
# =====================
# Gradio App
# =====================
def chat_fn(message, history, steps, max_len):
response = generate_with_prompt(model, message, max_length=max_len, T=steps)
history.append((message, response))
return "", history
with gr.Blocks() as demo:
gr.Markdown("## 🌀 DiffusionTextModel QA Chat Demo")
chatbot = gr.Chatbot()
msg = gr.Textbox(placeholder="Type your question or prompt here...")
steps = gr.Slider(1, 50, value=10, step=1, label="Diffusion Steps (T)")
max_len = gr.Slider(10, 99, value=50, step=1, label="Max Token Length (≤ 99)")
clear = gr.Button("Clear")
msg.submit(chat_fn, [msg, chatbot, steps, max_len], [msg, chatbot])
clear.click(lambda: None, None, chatbot, queue=False)
demo.launch()
|