Refactor start token handling in compute_entropy function

app.py

@@ -46,11 +46,16 @@ def compute_entropy(code: str):
         if attention_mask is not None:
             attention_mask = attention_mask.to(device)
 
-        # Prepend BOS if not already present so the first real token gets a predicted probability
-        bos_id = TOKENIZER.bos_token_id
-        if bos_id is not None and input_ids[0, 0].item() != bos_id:
-            bos_tensor = torch.full((1, 1), bos_id, dtype=input_ids.dtype, device=device)
-            input_ids = torch.cat([bos_tensor, input_ids], dim=1)
+        # Prepend a start token so the first real token gets a predicted probability.
+        # Prefer bos_token_id, fall back to eos/pad (pad is always set above).
+        start_id = next(
+            (getattr(TOKENIZER, a) for a in ("bos_token_id", "eos_token_id", "pad_token_id")
+             if getattr(TOKENIZER, a, None) is not None),
+            None,
+        )
+        if start_id is not None and input_ids[0, 0].item() != start_id:
+            start_tensor = torch.full((1, 1), start_id, dtype=input_ids.dtype, device=device)
+            input_ids = torch.cat([start_tensor, input_ids], dim=1)
             if attention_mask is not None:
                 attention_mask = torch.cat(
                     [torch.ones(1, 1, dtype=attention_mask.dtype, device=device), attention_mask], dim=1