Update README.md

README.md

@@ -9,8 +9,7 @@ tags:
 
 A 231M parameter base model trained on lichess games from January 2023 that ended in checkmate (filtered out games that were won because of time).
 
-It's trained on sequences of UCI moves, inference is pretty simple:
-
+## Inference
 ```py
 from transformers import GPT2LMHeadModel, AutoTokenizer
 
@@ -26,4 +25,18 @@ next_move = tokenizer.decode(gen_tokens[-1])
 print(next_move) #d4e5
 ```
 
+### End of game detection
+
+The model also has three special tokens for end game detection `<BLACK_WIN>`, `<WHITE_WIN>` and `<DRAW>`. This can be useful for implementing beam search strategies. 
+
+```py
+moves = " ".join(["f2f3", "e7e5", "g2g4", "d8h4"])
+
+model_inputs = tokenizer(moves, return_tensors="pt")
+gen_tokens = model.generate(**model_inputs, max_new_tokens=1)[0]
+next_move = tokenizer.decode(gen_tokens[-1])
+
+print(next_move) # <BLACK_WIN>
+```
+