philipp-zettl commited on
Commit
ee2b517
1 Parent(s): 9ede36f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -1,4 +1,5 @@
1
  import gradio as gr
 
2
  from model import DecoderTransformer, Tokenizer
3
  from huggingface_hub import hf_hub_download
4
  import torch
@@ -29,6 +30,7 @@ tokenizer = Tokenizer.from_pretrained(tokenizer_path)
29
  def generate(prompt):
30
  model_input = torch.tensor(tokenizer.encode(prompt), dtype=torch.long, device=device).view((1, len(prompt)))
31
  pgn = tokenizer.decode(model.generate(model_input, max_new_tokens=4, context_size=context_size)[0].tolist())
 
32
  game = chess.pgn.read_game(pgn)
33
  img = chess.svg.board(game.board())
34
  return pgn, img
 
1
  import gradio as gr
2
+ from io import StringIO
3
  from model import DecoderTransformer, Tokenizer
4
  from huggingface_hub import hf_hub_download
5
  import torch
 
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