yasserrmd's picture
Update app.py
4b51c4b verified
raw
history blame
2.3 kB
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()