Views
No views yet
1import torch
2import transformers
3from transformers import AutoTokenizer, AutoModelForCausalLM
4
5MIN_TRANSFORMERS_VERSION = '4.25.1'
6
7# check transformers version
8assert transformers.__version__ >= MIN_TRANSFORMERS_VERSION, f'Please upgrade transformers to version {MIN_TRANSFORMERS_VERSION} or higher.'
9
10# init
11tokenizer = AutoTokenizer.from_pretrained("Waterhorse/chessgpt-base-v1")
12model = AutoModelForCausalLM.from_pretrained("Waterhorse/chessgpt-base-v1", torch_dtype=torch.float16)
13model = model.to('cuda:0')
14
15# infer
16# Conversation between two
17prompt = "Q: 1.e4 c5, what is the name of this opening?A:"
18
19inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
20input_length = inputs.input_ids.shape[1]
21outputs = model.generate(
22 **inputs, max_new_tokens=128, do_sample=True, temperature=0.7, top_p=0.7, top_k=50, return_dict_in_generate=True,
23)
24token = outputs.sequences[0, input_length:]
25output_str = tokenizer.decode(token)
26print(output_str)chessgpt-base-v1 is mainly for research on large language model, especially for those research about policy learning and language modeling.chessgpt-base-v1 is a language model trained on chess related data and may not perform well for other use cases beyond chess domain.1@article{feng2023chessgpt,
2 title={ChessGPT: Bridging Policy Learning and Language Modeling},
3 author={Feng, Xidong and Luo, Yicheng and Wang, Ziyan and Tang, Hongrui and Yang, Mengyue and Shao, Kun and Mguni, David and Du, Yali and Wang, Jun},
4 journal={arXiv preprint arXiv:2306.09200},
5 year={2023}
6}