Change back to User-Assistant conversation

app.py

@@ -41,10 +41,10 @@ token_probabilities = np.array([token_probs_dict[str(i)] for i in range(len(toke
 
 def load_model():
     ckpt_path = hf_hub_download(
-        repo_id="ruurd/tini_bi",
+        repo_id="ruurd/tini",
         filename="diffusion-model.pth",
         token=os.getenv("HF_TOKEN"),
-        revision="5a22a8b6168466dbbf704efd00d8cbf2eee51426",
+        # revision="5a22a8b6168466dbbf704efd00d8cbf2eee51426",
     )
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -71,8 +71,6 @@ def load_model():
     model.eval()
     return model
 
-
-
 rng = np.random.default_rng()
 
 # --- Utility Functions ---
@@ -204,13 +202,11 @@ def diffusion_chat(question, eot_weight, mask_weight, max_it, pause_length, shar
 
     print('started generation')
     prompt = f"User: {question}\nAssistant:"
-    prompt = question
     input_ids = tokenizer.encode(prompt, add_special_tokens=False)
     answer_start = find_answer_start(input_ids, assistant_marker_ids)
-    # if answer_start is None:
-    #     yield "Error: Could not find Assistant marker in input."
-    #     return
-    answer_start = len(input_ids)
+    if answer_start is None:
+        yield "Error: Could not find Assistant marker in input."
+        return
     
     if len(input_ids) < 256:
         input_ids += [pad_token] * (256 - len(input_ids))