Fixed early-stopping in get_mamba_response based on space/dot tokens (now decodes the strings instead of using hardcoded token ids).

chess-gpt-eval/mamba_module.py

@@ -81,6 +81,8 @@ class MambaPlayer:
         self.vocab_size = vocab_size
         self.encode = encode
         self.decode = decode
+        self.space_tok = encode(' ')[0]
+        self.dot_tok = encode('.')[0]
         self.model = model
         self.ctx = ctx
         self.device = device
@@ -107,8 +109,9 @@ class MambaPlayer:
 
                 probs = torch.nn.functional.softmax(logits, dim=-1)
                 next_token_id = torch.multinomial(probs, num_samples=1)
-                if have_non_space and (next_token_id == 0 or next_token_id==4):
-                    break
+                if next_token_id == self.space_tok or next_token_id==self.dot_tok:
+                    if have_non_space:
+                        break
                 else:
                     have_non_space = True
                 input_ids = torch.cat([input_ids, next_token_id.unsqueeze(0)], dim=1)