Update app.py

app.py

@@ -70,13 +70,13 @@ def log_odds(p, eps=0):
 def make_pred_mask(input_ids):
     user_mask = [False] * len(input_ids)
     is_user_role = False
-    for i in range(len(input_ids)-1):
-        if input_ids[i] == START_TOKEN_ID:
-            is_user_role = input_ids[i+1] in USER_TOKEN_IDS
-        if is_user_role and (input_ids[i] not in SPECIAL_TOKENS):
-            user_mask[i] = True
-        else:
+    for i in range(len(input_ids)):
+        tok = input_ids[i]
+        if tok == START_TOKEN_ID and i + 1 < len(input_ids):
+            is_user_role = input_ids[i + 1] in USER_TOKEN_IDS
             user_mask[i] = False
+            continue
+        user_mask[i] = is_user_role and (tok not in SPECIAL_TOKENS)
     return user_mask