0.0.3 language hotfix

r4l/__init__.py

@@ -0,0 +1,10 @@
+lang_dict = {
+    'en': ['tts_models/en/ljspeech/tacotron2-DDC', 'vocoder_models/en/ljspeech/hifigan_v2', 'eng'],
+    'fr': ['tts_models/fr/mai/tacotron2-DDC', 'vocoder_models/universal/libri-tts/fullband-melgan', 'fra'],
+    'es': ['tts_models/es/mai/tacotron2-DDC', 'vocoder_models/universal/libri-tts/fullband-melgan', 'spa'],
+    'de': ['tts_models/de/thorsten/tacotron2-DCA', 'vocoder_models/de/thorsten/fullband-melgan', 'deu'],
+    'ja': ['tts_models/ja/kokoro/tacotron2-DDC', 'vocoder_models/ja/kokoro/hifigan_v1', 'jpn'],
+    'nl': ['tts_models/nl/mai/tacotron2-DDC', 'vocoder_models/nl/mai/parallel-wavegan', 'nld'],
+    'zh_sim': ['tts_models/zh-CN/baker/tacotron2-DDC-GST', None, 'chi_sim'],
+    'zh_tra': ['tts_models/zh-CN/baker/tacotron2-DDC-GST', None, 'chi_tra']
+}

r4l/bin/cli.py

@@ -4,7 +4,8 @@
 import sys
 
 from r4l.util.text import TextProcessor
-from r4l.util.reader import Reader, models_dict
+from r4l.util.reader import Reader
+from r4l import lang_dict
 import os
 import time
 import csv
@@ -28,7 +29,7 @@ def get_ext(filename):
 
 def get_texts(sesspath, lang, force_english):
     wordcount = 0
-    tp = TextProcessor(sc_langs=lang)
+    tp = TextProcessor(langs=lang)
     files = [f for f in os.listdir(sesspath) if get_ext(f) in ['pdf', 'txt', 'muse']]
     texts = [[] for _ in files]
     print(f"> Reading {files}")
@@ -86,7 +87,7 @@ def main():
     )
     args = parser.parse_args()
     if args.list_langs:
-        print(models_dict.keys())
+        print(list(lang_dict.keys()))
         sys.exit()
     if not os.path.isdir(args.in_path):
         print("input path must exist and contain files!")

r4l/util/reader.py

@@ -4,16 +4,7 @@
 import numpy as np
 from pathlib import Path
 from r4l.util.text import split_into_sentences
-
-models_dict = {
-    'en': ['tts_models/en/ljspeech/tacotron2-DDC', 'vocoder_models/en/ljspeech/hifigan_v2'],
-    'fr': ['tts_models/fr/mai/tacotron2-DDC', 'vocoder_models/universal/libri-tts/fullband-melgan'],
-    'es': ['tts_models/es/mai/tacotron2-DDC', 'vocoder_models/universal/libri-tts/fullband-melgan'],
-    'de': ['tts_models/de/thorsten/tacotron2-DCA', 'vocoder_models/de/thorsten/fullband-melgan'],
-    'ja': ['tts_models/ja/kokoro/tacotron2-DDC', 'vocoder_models/ja/kokoro/hifigan_v1'],
-    'nl': ['tts_models/nl/mai/tacotron2-DDC', 'vocoder_models/nl/mai/parallel-wavegan'],
-    'zh': ['tts_models/zh-CN/baker/tacotron2-DDC-GST', None]
-}
+from r4l import lang_dict
 
 # later i'll figure out how to load TTS's .models.json
 
@@ -22,7 +13,7 @@
 class Reader:
     def __init__(self, outpath, lang='en', tts_name=None, voc_name=None):
         self.outpath = outpath
-        model_name, vocoder_name = models_dict[lang]
+        model_name, vocoder_name, _ = lang_dict[lang]
         if tts_name is not None:
             model_name = tts_name
         if voc_name is not None:

r4l/util/text.py

@@ -7,6 +7,7 @@
 from difflib import SequenceMatcher
 from ocrmypdf import ocr
 import os
+from r4l import lang_dict
 
 def only_english(text):
     import nltk
@@ -38,12 +39,12 @@ def split_into_sentences(string):
 spec = dict((re.escape(k), v) for k, v in spec.items())
 spec_re = re.compile("|".join(spec.keys()))
 
-
 class TextProcessor:
-    def __init__(self, bert_model="distilbert-base-multilingual-cased", sc_langs=["en", "fr"]):
+    def __init__(self, bert_model="distilbert-base-multilingual-cased", langs=["en", "fr"]):
         self.tokenizer = AutoTokenizer.from_pretrained(bert_model)
         self.model = AutoModelForMaskedLM.from_pretrained(bert_model)
-        self.sc = SpellChecker(distance=1, language=sc_langs)
+        self.sc = SpellChecker(distance=1, language=langs)
+        self.lang = [l[3] for l in lang_dict[langs]]
         print("> BERT initialized")
 
     # get and correct text
@@ -69,7 +70,7 @@ def _load(self, filename, sesspath, force):
         if not os.path.isdir(sesspath + "tmp/"):
             os.mkdir(sesspath + "tmp/")
         ocr(sesspath + filename, sesspath + "tmp/tmp.pdf", sidecar=tpath, redo_ocr=(not force), deskew=force,
-            rotate_pages=force, remove_background=force, clean=force, force_ocr=force)
+            rotate_pages=force, remove_background=force, clean=force, force_ocr=force,language=self.lang)
         with open(tpath, "r") as txt:
             text = txt.read()
         print("> OCR complete")

setup.py

@@ -7,7 +7,7 @@
     readme = f.read()
 setup(
     name='reading4listeners',
-    version='0.0.2',
+    version='0.0.3',
     packages=['r4l'],
     url='https://github.com/CypherousSkies/reading-for-listeners',
     project_urls={