Update README.md

README.md

@@ -40,22 +40,24 @@ Two possible simple extensions:
 import chess
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+
 def next_move(model, tokenizer, fen):
     input_ids = tokenizer(f"FEN: {fen}\nMOVE:", return_tensors="pt")
-    input_ids = {k:v.to(model.device) for k,v in input_ids.items()}
+    input_ids = {k: v.to(model.device) for k, v in input_ids.items()}
     out = model.generate(
-        **input_ids, 
+        **input_ids,
         max_new_tokens=10,
         pad_token_id=tokenizer.eos_token_id,
         do_sample=True,
         temperature=0.1,
-        )
+    )
     out_str = tokenizer.batch_decode(out)[0]
-    return out_str.split('MOVE:')[-1].replace("<|endoftext|>", "").strip()
+    return out_str.split("MOVE:")[-1].replace("<|endoftext|>", "").strip()
+
 
 board = chess.Board()
-model = AutoModelForCausalLM.from_pretrained('Xmaster6y/gpt2-stockfish-debug')
-tokenizer = AutoTokenizer.from_pretrained('gpt2')
+model = AutoModelForCausalLM.from_pretrained("Xmaster6y/gpt2-stockfish-debug")
+tokenizer = AutoTokenizer.from_pretrained("Xmaster6y/gpt2-stockfish-debug")  # or "gpt2"
 tokenizer.pad_token = tokenizer.eos_token
 for i in range(100):
     fen = board.fen()