Bug with OPUS model

helpers/batching.py

@@ -435,11 +435,13 @@ def generate_text(
     return all_generated_texts
 
 if __name__ == '__main__':
-    GROQ_API_KEY = os.getenv('GROQ_API_KEY')
-    GEMINI_API_KEY = os.getenv('GEMINI_API_KEY')
-    groq = Groqq('gemma-7b-it', 15, GROQ_API_KEY)
-    groq.set_lang('zh','en')
-    gemini = Gemini('gemini-1.5-pro', 15, GEMINI_API_KEY)
-    gemini.set_lang('zh','en')
-    print(gemini.translate(['荷兰咯']))
-    print(groq.translate(['荷兰咯']))
+    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, M2M100Tokenizer, M2M100ForConditionalGeneration
+    opus_model = 'Helsinki-NLP/opus-mt-en-zh'
+    LOCAL_FILES_ONLY = True
+    tokenizer = AutoTokenizer.from_pretrained(opus_model, local_files_only=LOCAL_FILES_ONLY)
+    model = AutoModelForSeq2SeqLM.from_pretrained(opus_model, local_files_only=LOCAL_FILES_ONLY, torch_dtype=torch.float16).to(device)
+    # tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", local_files_only=LOCAL_FILES_ONLY)
+    # tokenizer.src_lang = "en"
+    # model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M", local_files_only=LOCAL_FILES_ONLY).to(device)
+    
+    print(generate_text([ i.lower().capitalize() for i in ['placeholder','Story','StoRY', 'TufoRIaL', 'CovFfG', 'LoaD DaTA', 'SAME DATa', 'ReTulN@TitIE', 'View', '@niirm', 'SysceM', 'MeNu:', 'MaND', 'CoM', 'SeLEcT', 'Frogguingang', 'Tutorias', 'Back']], model, tokenizer))

helpers/ocr.py

@@ -87,11 +87,12 @@ def _id_filtered(ocr, image, lang) -> list:
 def _id_lang(ocr, image, lang) -> list:
     result = _identify(ocr, image)
     lang = standardize_lang(lang)['id_model_lang']
-    try:
-        filtered = [entry for entry in result if contains_lang(entry[1], lang)]
-    except:
-        logger.error(f"Selected language not part of default: {default_languages}.")
-        raise ValueError(f"Selected language not part of default: {default_languages}.")
+    # try:
+    logger.info(f"Filtering out phrases not in {lang}.")
+    filtered = [entry for entry in result if contains_lang(entry[1], lang)]
+    # except:
+        # logger.error(f"Selected language not part of default: {default_languages}.")
+        # raise ValueError(f"Selected language not part of default: {default_languages}.")
     return filtered
 
 def id_keep_source_lang(ocr, image, lang) -> list:

helpers/translation.py

@@ -143,15 +143,17 @@ def get_OPUS_model(from_lang, target_lang):
 
 def init_OPUS(from_lang = 'ch_sim', target_lang = 'en'):
     opus_model = get_OPUS_model(from_lang, target_lang)
+    logger.debug(f"OPUS model: {opus_model}")
     tokenizer = AutoTokenizer.from_pretrained(opus_model, local_files_only=LOCAL_FILES_ONLY)
     model = AutoModelForSeq2SeqLM.from_pretrained(opus_model, local_files_only=LOCAL_FILES_ONLY, torch_dtype=torch.float16).to(device)
     model.eval()
     return (model, tokenizer)
 
 def translate_OPUS(text: list[str], model, tokenizer) -> list[str]:
-    translated_text = generate_text(model,tokenizer, text, 
+    translated_text = generate_text(text, model,tokenizer, 
                   batch_size=BATCH_SIZE, device=device, 
                   max_length=MAX_INPUT_TOKENS, max_new_tokens=MAX_OUTPUT_TOKENS)
+    logger.debug(f"Translated text: {translated_text}")
     return translated_text
 
 ###############################
@@ -174,6 +176,7 @@ def translate_Seq_LLM(text,
                         model,
                         tokenizer,
                        **kwargs):
+    text = [t.lower().capitalize() for t in text]
     if model_type == 'opus':
         return translate_OPUS(text, model, tokenizer)
     elif model_type == 'm2m':

helpers/utils.py

@@ -28,7 +28,7 @@ def is_wayland():
 
 #  please install grim otherwise this is way too slow for wayland
 def printsc_wayland(region: tuple, path: str):
-    subprocess.run(['grim','-g', f'{region[0]},{region[1]} {region[2]-region[0]}x{region[3]-region[1]}', '-t', 'jpeg', '-q','100', path])
+    subprocess.run(['grim','-g', f'{region[0]},{region[1]} {region[2]-region[0]}x{region[3]-region[1]}', '-t', 'jpeg', '-q','90', path])
 
 def printsc_non_wayland(region: tuple, path: str):
     # use mss to capture the screen