Update.

[pytorch.git] / gpt-test.py

diff --git a/gpt-test.py b/gpt-test.py
index ff72e50..557f734 100755 (executable)
--- a/gpt-test.py
+++ b/gpt-test.py
@@ -22,7 +22,9 @@ from transformers import GPT2Tokenizer, GPT2LMHeadModel
 
 ######################################################################
 
-def complete(model, primer, nb_sentences = 1, nb_token_max = 100, temperature = None):
+def complete(model, tokenizer,
+             primer,
+             nb_sentences = 1, nb_token_max = 100, temperature = None):
     nt, ns = 0, 0
     tokens = tokenizer.encode(primer)
     primer_len = len(tokens)
@@ -52,7 +54,7 @@ model = GPT2LMHeadModel.from_pretrained(model_name)
 model.eval()
 
 print(
-    complete(model,
+    complete(model, tokenizer,
              'The object was blue all over, but also green all over, it was a',
     )
 )