Silas Blanchard commited on
Commit
ed87abf
·
unverified ·
1 Parent(s): 7112478

Update model-runner.py

Browse files
Files changed (1) hide show
  1. src/model-runner.py +4 -1
src/model-runner.py CHANGED
@@ -1,5 +1,8 @@
 
1
  import chess
2
- import random
 
 
3
 
4
  def cleanup_output(text, prompt, extra_len):
5
  section = text[len(prompt):]
 
1
+ from transformers import pipeline
2
  import chess
3
+
4
+ generator2 = pipeline('text-generation', model='BlueSunflower/gpt2-medium-chess')
5
+ generator = pipeline('text-generation', model='gpt2')
6
 
7
  def cleanup_output(text, prompt, extra_len):
8
  section = text[len(prompt):]