projects
/
pytorch.git
/ blobdiff
commit
grep
author
committer
pickaxe
?
search:
re
summary
|
shortlog
|
log
|
commit
|
commitdiff
|
tree
raw
|
inline
| side by side
Update
[pytorch.git]
/
gpt-test.py
diff --git
a/gpt-test.py
b/gpt-test.py
index
0967043
..
ddd7dcf
100755
(executable)
--- a/
gpt-test.py
+++ b/
gpt-test.py
@@
-22,9
+22,10
@@
from transformers import GPT2Tokenizer, GPT2LMHeadModel
######################################################################
######################################################################
-def complete(model, tokenizer,
- 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)
nt, ns = 0, 0
tokens = tokenizer.encode(primer)
primer_len = len(tokens)
@@
-33,31
+34,41
@@
def complete(model, tokenizer,
if temperature is None:
next_token = torch.argmax(outputs[0, -1])
else:
if temperature is None:
next_token = torch.argmax(outputs[0, -1])
else:
- dist =
torch.distributions.Categorical(logits =
outputs[0, -1] / temperature)
+ dist =
torch.distributions.Categorical(logits=
outputs[0, -1] / temperature)
next_token = dist.sample((1,)).item()
tokens.append(next_token)
nt += 1
next_token = dist.sample((1,)).item()
tokens.append(next_token)
nt += 1
- if tokenizer.decode([next_token]) == '.': ns += 1
+ if tokenizer.decode([next_token]) == ".":
+ ns += 1
if ns == nb_sentences or nt == nb_token_max:
if ns == nb_sentences or nt == nb_token_max:
- return '<' + tokenizer.decode(tokens[:primer_len]) + '>' + \
- tokenizer.decode(tokens[primer_len:])
+ return (
+ "<"
+ + tokenizer.decode(tokens[:primer_len])
+ + ">"
+ + tokenizer.decode(tokens[primer_len:])
+ )
+
######################################################################
######################################################################
-#model_name = 'gpt2'
-#model_name = 'gpt2-large'
-model_name =
'gpt2-xl'
+#
model_name = 'gpt2'
+#
model_name = 'gpt2-large'
+model_name =
"gpt2-xl"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
model.eval()
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
model.eval()
-print(f'Using {model_name} ({int(sum(p.numel() for p in model.parameters())/(1e6))}M parameters)')
+print(
+ f"Using {model_name} ({int(sum(p.numel() for p in model.parameters())/(1e6))}M parameters)"
+)
print(
print(
- complete(model, tokenizer,
- 'The object was blue all over, but also green all over, it was a',
+ complete(
+ model,
+ tokenizer,
+ "The object was blue all over, but also green all over, it was a",
)
)
)
)