philipp-zettl commited on
Commit
22bfac3
1 Parent(s): ee2b517

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -30,10 +30,13 @@ tokenizer = Tokenizer.from_pretrained(tokenizer_path)
30
  def generate(prompt):
31
  model_input = torch.tensor(tokenizer.encode(prompt), dtype=torch.long, device=device).view((1, len(prompt)))
32
  pgn = tokenizer.decode(model.generate(model_input, max_new_tokens=4, context_size=context_size)[0].tolist())
33
- pgn = StringIO(pgn)
34
- game = chess.pgn.read_game(pgn)
35
  img = chess.svg.board(game.board())
36
- return pgn, img
 
 
 
37
 
38
 
39
  with gr.Blocks() as demo:
 
30
  def generate(prompt):
31
  model_input = torch.tensor(tokenizer.encode(prompt), dtype=torch.long, device=device).view((1, len(prompt)))
32
  pgn = tokenizer.decode(model.generate(model_input, max_new_tokens=4, context_size=context_size)[0].tolist())
33
+ pgn_str = StringIO(pgn)
34
+ game = chess.pgn.read_game(pgn_str)
35
  img = chess.svg.board(game.board())
36
+ filename = f'moves-{pgn}'
37
+ with open(filename, 'w') as f:
38
+ f.write(img)
39
+ return pgn, filename
40
 
41
 
42
  with gr.Blocks() as demo: