From cd0426ab9ea96032e7ed9ad3b728e6ff51977c7e Mon Sep 17 00:00:00 2001 From: Anthony Wang Date: Tue, 22 Feb 2022 22:25:03 -0600 Subject: Pick a random initial input if none provided --- bot.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/bot.py b/bot.py index ea86b88..d9d0d93 100644 --- a/bot.py +++ b/bot.py @@ -1,4 +1,5 @@ from argparse import ArgumentParser +from random import choice from mastodon import Mastodon from transformers import AutoTokenizer, AutoModelForCausalLM @@ -17,10 +18,28 @@ tokenizer = AutoTokenizer.from_pretrained('distilgpt2') model = AutoModelForCausalLM.from_pretrained(args.model) +if args.input is None: + # Create random input + args.input = choice([ + 'I am', + 'My life is', + 'Computers are', + 'This is', + 'My', + 'I\'ve', + 'No one', + 'I love', + 'I will die of', + 'I', + 'The', + 'Anime' + ]) + + # Run the input through the model inputs = tokenizer.encode(args.input, return_tensors="pt") output = tokenizer.decode(model.generate( - inputs, do_sample=True, max_length=75, top_p=0.9)[0]) + inputs, do_sample=True, max_length=100, top_p=0.9)[0]) print(output) @@ -32,4 +51,4 @@ mastodon = Mastodon( post = output.split('\n')[0] if len(post) < 100: post = output.split('\n')[0] + '\n' + output.split('\n')[1] -mastodon.status_post(post) +mastodon.status_post(post[:500]) -- cgit v1.2.3-70-g09d2