diff options
| -rw-r--r-- | bot.py | 19 |
1 files changed, 10 insertions, 9 deletions
@@ -1,5 +1,6 @@ from argparse import ArgumentParser from random import randint, choice +from re import sub from transformers import AutoTokenizer, AutoModelForCausalLM @@ -61,11 +62,6 @@ if args.input is None: while len(line) < 2: line = choice(lines).split() - # Remove mentions - if line[0].count('@') > 1: - line[0] = '@'.join(line[0].split('@')[0:2]) - if line[1].count('@') > 1: - line[1] = '@'.join(line[1].split('@')[0:2]) args.input = line[0] + ' ' + line[1] @@ -77,10 +73,6 @@ output = tokenizer.decode(model.generate( print(output) -if args.instance is None: - quit() - - # Prepare the post output = output.split('\n') post = output[0] @@ -89,6 +81,15 @@ if len(post) < 200 and len(output) > 1: post = post[:500] +# Remove mentions +post = sub('(@[^ ]*)@[^ ]*', '\\1', post) + + +# Quit if no instance specified +if args.instance is None: + quit() + + # Post it! if args.backend == 'mastodon': from mastodon import Mastodon |