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

1 files changed, 21 insertions, 1 deletions

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)