File size: 1,187 Bytes
eada88b ab13c7b eada88b e9fc112 eada88b e9fc112 eada88b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
---
license: mit
datasets:
- nsarrazin/lichess-games-2023-01
pipeline_tag: text-generation
tags:
- chess
---
A 231M parameter base model trained on 4.4B tokens of lichess games from January 2023 that ended in checkmate (filtered out games that were won because of time).
## Inference
```py
from transformers import GPT2LMHeadModel, AutoTokenizer
model = GPT2LMHeadModel.from_pretrained("nsarrazin/chessformer").eval()
tokenizer = AutoTokenizer.from_pretrained("nsarrazin/chessformer")
moves = " ".join(["e2e4", "e7e5", "d2d4", "d7d5"])
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) #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>
```
|