Pick a random initial input if none provided

1 files 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])