Actually remove mentions

1 files changed, 10 insertions, 9 deletions

diff --git a/bot.py b/bot.py
index f392b11..0726fb5 100644
--- a/bot.py
+++ b/bot.py
@@ -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