Upload model

chat_utils.py

@@ -11,20 +11,7 @@ import torch
 import torch.nn.functional as F
 from transformers import AutoTokenizer
 
-sys.path.insert(1, "/lustre/fsw/portfolios/nvr/users/yongganf/adlr-megatron-lm")
-from get_hf_model import get_torchtitan_model_sft  # noqa: E402
 
-
-# --------------------------- Reproducibility ----------------------------------
-def set_seed(seed: int = 42):
-    torch.manual_seed(seed)
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-
-# -------------------- Diffusion helpers (unchanged logic) --------------------
 def get_transfer_index(
     logits, temperature, remasking, mask_index, x, num_transfer_tokens, threshold=None, neg_entropy=False
 ):

modeling_nvrdiff.py

@@ -535,41 +535,18 @@ class DiffEncoderModel(Qwen3PreTrainedModel, GenerationMixin):
         )
     
 
-    def chat(self, tokenizer, max_new_tokens, steps, block_length, threshold):
-        print("Stateless chat (type 'exit' to quit)")
-        print("------------------------------------")
-
-        try:
-            while True:
-                user_input = input("User: ").strip()
-                if user_input.lower() in {"exit", "quit", "q"}:
-                    print("Conversation ended.")
-                    break
-
-                prompt_ids = tokenizer(
-                    user_input,return_tensors='pt'
-                ).input_ids.to(device='cuda')
-
-                out_ids, nfe = generate_with_prefix_cache_block_diff(
-                    model=self,
-                    prompt=prompt_ids,
-                    gen_length=max_new_tokens,
-                    steps=steps,
-                    block_length=block_length,
-                    remasking="low_confidence",
-                    mask_id=self.mask_token_id,
-                    threshold=threshold,
-                    shift_logits=True,
-                    neg_entropy=True,
-                )
-
-                generated_tokens = out_ids[:, prompt_ids.shape[1]:]
-                tokenized_out = tokenizer.batch_decode(
-                    generated_tokens,
-                    skip_special_tokens=True
-                )[0]
-                print(f"Model: {tokenized_out}")
-                print(f"[nfe={nfe}]")
-
-        except KeyboardInterrupt:
-            print("\n[info] interrupted by user (Ctrl-C).")
+    def generate(self, prompt_ids, max_new_tokens, steps, block_length, threshold):
+        out_ids, nfe = generate_with_prefix_cache_block_diff(
+                        model=self,
+                        prompt=prompt_ids,
+                        gen_length=max_new_tokens,
+                        steps=steps,
+                        block_length=block_length,
+                        remasking="low_confidence",
+                        mask_id=self.mask_token_id,
+                        threshold=threshold,
+                        shift_logits=True,
+                        neg_entropy=True,
+                    )
+                    
+        return out_ids, nfe