feat: spinner

app.py

@@ -72,39 +72,41 @@ diffusion = get_diffusion()
 
 sample_fn = diffusion.ddim_sample_loop
 
+st.title("Lang2mol-Diff")
 text_input = st.text_area("Enter molecule description")
-output = tokenizer(
-    text_input,
-    max_length=256,
-    truncation=True,
-    padding="max_length",
-    add_special_tokens=True,
-    return_tensors="pt",
-    return_attention_mask=True,
-)
-caption_state = encoder(
-    input_ids=output["input_ids"],
-    attention_mask=output["attention_mask"],
-).last_hidden_state
-caption_mask = output["attention_mask"]
+with st.spinner("Please wait ..."):
+    output = tokenizer(
+        text_input,
+        max_length=256,
+        truncation=True,
+        padding="max_length",
+        add_special_tokens=True,
+        return_tensors="pt",
+        return_attention_mask=True,
+    )
+    caption_state = encoder(
+        input_ids=output["input_ids"],
+        attention_mask=output["attention_mask"],
+    ).last_hidden_state
+    caption_mask = output["attention_mask"]
 
-outputs = sample_fn(
-    model,
-    (1, 256, 32),
-    clip_denoised=False,
-    denoised_fn=None,
-    model_kwargs={},
-    top_p=1.0,
-    progress=True,
-    caption=(caption_state, caption_mask),
-)
-logits = model.get_logits(torch.tensor(outputs))
-cands = torch.topk(logits, k=1, dim=-1)
-outputs = cands.indices
-outputs = outputs.squeeze(-1)
-outputs = tokenizer.decode(outputs)
-result = sf.decoder(
-    outputs[0].replace("<pad>", "").replace("</s>", "").replace("\t", "")
-).replace("\t", "")
+    outputs = sample_fn(
+        model,
+        (1, 256, 32),
+        clip_denoised=False,
+        denoised_fn=None,
+        model_kwargs={},
+        top_p=1.0,
+        progress=True,
+        caption=(caption_state, caption_mask),
+    )
+    logits = model.get_logits(torch.tensor(outputs))
+    cands = torch.topk(logits, k=1, dim=-1)
+    outputs = cands.indices
+    outputs = outputs.squeeze(-1)
+    outputs = tokenizer.decode(outputs)
+    result = sf.decoder(
+        outputs[0].replace("<pad>", "").replace("</s>", "").replace("\t", "")
+    ).replace("\t", "")
 
-st.write(result)
+    st.write(result)