Fix loading of model and tokenizer

Change location of inference functions

app.py

@@ -29,56 +29,8 @@ hf_token = os.getenv("HF_TOKEN")
 if hf_token is None:
     raise ValueError("HF_TOKEN is not set")
 
-# --- Load tokenizer ---
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B", use_fast=True, token=hf_token)
-vocab_size = len(tokenizer)
-eos_token_id = tokenizer.eos_token_id
-mask_token_id = tokenizer.encode('MASK', add_special_tokens=False)[0]
-assistant_marker_ids = tokenizer.encode("<|start_header_id|>assistant<|end_header_id|>", add_special_tokens=False)
-
 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, clustering=0.5, noise_start = 1.0):
-    noised = input_ids.copy()
-    answer_len = len(noised) - answer_start
-    num_to_noise = int(threshold * answer_len * noise_start)
-    mask_token_id = tokenizer.encode('MASK', add_special_tokens = False)[0]
-
-    if num_to_noise == 0:
-        return noised, []
-
-    num_clusters = max(1, int((1 - clustering) * num_to_noise))
-    cluster_size = max(1, int(num_to_noise / num_clusters))
-
-    noised_indices = set()
-    for _ in range(num_clusters):
-        center = rng.integers(answer_start, len(noised))
-        span_start = max(answer_start, center - cluster_size // 2)
-        span_end = min(len(noised), span_start + cluster_size)
-        noised_indices.update(range(span_start, span_end))
-
-    noised_indices = sorted(list(noised_indices))[:num_to_noise]
-
-    for idx in noised_indices:
-        noised[idx] = mask_token_id
-
-    return noised, noised_indices
-
-
 # Add new noising function
 def confidence_guided_noising(input_ids, answer_start, confidences, noise_clipping, threshold=1.0, noise_start=1.0):
     noised = input_ids.copy()
@@ -121,33 +73,6 @@ def confidence_guided_noising(input_ids, answer_start, confidences, noise_clippi
     noised_indices = sorted(noised_indices)
     return noised, noised_indices
 
-def filter_logits(logits, top_k=0, top_p=0.0):
-    """Filter logits per position for top-k / nucleus (top-p) sampling."""
-    logits = logits.clone()  # don't modify in-place
-    batch_size, seq_len, vocab_size = logits.shape
-
-    for i in range(seq_len):
-        token_logits = logits[0, i]
-
-        if top_k > 0:
-            top_values, _ = torch.topk(token_logits, top_k)
-            threshold = top_values[-1]
-            token_logits[token_logits < threshold] = float("-inf")
-
-        if top_p > 0.0:
-            sorted_logits, sorted_indices = torch.sort(token_logits, descending=True)
-            cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1)
-
-            sorted_indices_to_remove = cumulative_probs > top_p
-            sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
-            sorted_indices_to_remove[0] = 0  # always keep at least 1 token
-
-            token_logits[sorted_indices[sorted_indices_to_remove]] = float("-inf")
-
-        logits[0, i] = token_logits
-
-    return logits
-
 @spaces.GPU
 def generate_diffusion_text(input_ids, top_p, top_k):
     with torch.no_grad():
@@ -198,7 +123,7 @@ def diffusion_chat(question, max_it, pause_length, sharpness,
 
     ori_input_tokens = input_ids
     current_tokens, just_noised_indices = noisify_answer(
-                input_ids, answer_start, threshold=1.0, clustering=clustering, noise_start = 1.0,
+                input_ids, answer_start, tokenizer, threshold=1.0, clustering=clustering, noise_start = 1.0,
             )
     yield f"<b>Iteration 0 (initial noise):</b><br>" + tokenizer.decode(current_tokens[answer_start:], skip_special_tokens=True).replace('\n', '<br>')
     time.sleep(pause_length)
@@ -257,7 +182,7 @@ def diffusion_chat(question, max_it, pause_length, sharpness,
             # just_noised_indices = []
         else:
             noised_answer, just_noised_indices = noisify_answer(
-                current_tokens, answer_start, threshold=threshold, clustering=clustering, noise_start = noise_start,
+                current_tokens, answer_start, tokenizer, threshold=threshold, clustering=clustering, noise_start = noise_start,
             )
 
         # --- RED HIGHLIGHT ---
@@ -302,9 +227,14 @@ ckpt_path = hf_hub_download(
     filename="diffusion-model.pth",
     token=os.getenv("HF_TOKEN")
 )
-model = load_trained_model(checkpoint_path=ckpt_path)
+model, tokenizer = load_trained_model(checkpoint_path=ckpt_path)
 print("✅ Model loaded.")
 
+vocab_size = len(tokenizer)
+eos_token_id = tokenizer.eos_token_id
+mask_token_id = tokenizer.encode('MASK', add_special_tokens=False)[0]
+assistant_marker_ids = tokenizer.encode("<|start_header_id|>assistant<|end_header_id|>", add_special_tokens=False)
+
 demo = gr.Interface(
     fn=diffusion_chat,
     inputs=[

infer.py

@@ -82,8 +82,8 @@ def filter_logits(logits, top_k=0, top_p=1.0, temperature=1.0):
 
     return logits
 
-
-def decode_tokens_safe(tokenizer, token_ids):
+# --- Utility Functions ---
+def decode_tokens_safe(token_ids, tokenizer):
     return tokenizer.decode(token_ids, skip_special_tokens=True).replace("\n", " ")
 
 def find_answer_start(input_ids, marker_ids):
@@ -92,24 +92,36 @@ def find_answer_start(input_ids, marker_ids):
             return i + len(marker_ids)
     return None
 
-def noisify_answer(input_ids, answer_start, threshold=1.0, is_unmasked=None, mask_token_id=128002):
-    noised = input_ids.copy()
-    total_len = len(input_ids)
-    candidates = [
-        i for i in range(answer_start, total_len)
-        if is_unmasked is None or not is_unmasked[i]
-    ]
-    num_to_add = int(threshold * total_len)
-    if num_to_add > 0 and len(candidates) > 0:
-        newly_masked = rng.choice(candidates, size=min(num_to_add, len(candidates)), replace=False)
-        for idx in newly_masked:
-            noised[idx] = mask_token_id
-    return noised
-
 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, tokenizer, threshold=1.0, clustering=0.5, noise_start = 1.0):
+    noised = input_ids.copy()
+    answer_len = len(noised) - answer_start
+    num_to_noise = int(threshold * answer_len * noise_start)
+    mask_token_id = tokenizer.encode('MASK', add_special_tokens = False)[0]
+
+    if num_to_noise == 0:
+        return noised, []
+
+    num_clusters = max(1, int((1 - clustering) * num_to_noise))
+    cluster_size = max(1, int(num_to_noise / num_clusters))
+
+    noised_indices = set()
+    for _ in range(num_clusters):
+        center = rng.integers(answer_start, len(noised))
+        span_start = max(answer_start, center - cluster_size // 2)
+        span_end = min(len(noised), span_start + cluster_size)
+        noised_indices.update(range(span_start, span_end))
+
+    noised_indices = sorted(list(noised_indices))[:num_to_noise]
+
+    for idx in noised_indices:
+        noised[idx] = mask_token_id
+
+    return noised, noised_indices
+
 import torch.nn.functional as F
 
 def generate_diffusion_text(model, input_ids, answer_start, top_k=0, top_p=1.0, temperature=1.0,