Preserve and re-add emojis from the original data

Could I have done this by messing with the markov
chain directly? maybe. I don't care enough to try

3 files changed, 36 insertions, 2 deletions

diff --git a/config__.default.py b/config__.default.py
index a4a7a15..2a19083 100644
--- a/config__.default.py
+++ b/config__.default.py
@@ -12,4 +12,13 @@ character_limit = 80
 
 log_length = 168
 # How many of the most recently generated posts to log, so that the bot won't repeat them
-# The default value of 168 logs a week of hourly posts
\ No newline at end of file
+# The default value of 168 logs a week of hourly posts
+
+emojis = True
+# The markov chain strips colons, which removes emojis from posts but leaves their names.
+# If set to True, a list of emoji shortcodes (eg. :neofox_floof:) present in the dataset
+# will be stored and any occurences of the emojis' names will be replaced with the code to
+# add the emoji back in. This may cause problems with emojis whose names are common words,
+# eg. :cat:
+# If set to false, emoji names will be removed as well.
+# Reimport after changing this setting.
\ No newline at end of file
diff --git a/generate.py b/generate.py
index ee8cc3a..0102455 100644
--- a/generate.py
+++ b/generate.py
@@ -44,6 +44,11 @@ posts_log.append(text)
 with open(log_path, "w") as outfile:
   outfile.write(json.dumps(posts_log[-config__.log_length:]))
 
+if(config__.emojis == True):
+    emojis = json.load(open("emojis.json"))
+    for emoji in emojis:
+        text = text.replace(emoji, ":"+emoji+":")
+
 text = text.replace('@','@​').replace('#','#​')
 print(text)
 
diff --git a/import-misskey.py b/import-misskey.py
index 7e199c3..29b63bb 100644
--- a/import-misskey.py
+++ b/import-misskey.py
@@ -1,6 +1,7 @@
 import re
 import json
 import sys
+import config__
 
 print('[+] loading nlp enhanced markov chain')
 import chain
@@ -11,6 +12,7 @@ MFM_END = re.compile(r'\]+')
 HTML = re.compile(r'</?[a-z]+>')
 SPACE = re.compile(r'[ \n]+')
 CONTRACTION = re.compile(r"(\w+)'(\w+)")
+EMOJI = re.compile(r':[a-zA-Z0-9-_]+:')
 
 print('[+] loading note json')
 export_f = sys.argv[1]
@@ -18,6 +20,7 @@ export = open(export_f)
 export_json = json.load(export)
 
 corpus = []
+emojis = []
 
 for note in export_json:
     if note.get('visibility') not in ['public', 'unlisted']:
@@ -40,15 +43,32 @@ for note in export_json:
     text = re.sub(HTML, '', text)
     text = re.sub(SPACE, ' ', text)
     text = re.sub(CONTRACTION, r'\1\2', text)
+    if(config__.emojis == False):
+        text = re.sub(EMOJI, '', text)
     text = text.strip()
 
     print(f"     - {text}")
     corpus.append(text)
 
+    if(config__.emojis == True):
+        noteemojis = EMOJI.findall(text)
+        if(noteemojis):
+            emojis = emojis + noteemojis
+
+if(config__.emojis == True):
+    print('[+] preserving emojis:')
+    emojis = list(set(emojis))
+    s_emo = []
+    for emo in emojis:
+        s_emo.append(emo.replace(":", ""))
+    print(s_emo)
+    with open("emojis.json", "w") as f:
+        json.dump(s_emo, f)
+
 print('[+] building markov chain')
 model = chain.Text("@@note@@".join(corpus), well_formed=False)
 model_json = model.compile().to_json()
 
 print('[+] exporting')
-export = open(export_f.replace('.json', '.model.json'), 'w')
+export = open('model.json', 'w')
 export.write(model_json)