Added GEMINI API support, added proper support for Japanese romanisation and furigana and optimised batching for local LLMs

helpers/batching.py

@@ -0,0 +1,182 @@
+from torch.utils.data import Dataset, DataLoader
+from typing import List, Dict
+from dotenv import load_dotenv
+import os , sys, torch, time
+from multiprocessing import Process, Event
+load_dotenv()
+sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+from config import device
+from logging_config import logger
+
+
+class Gemini():
+    def __init__(self, name, rate):
+        self.name = name 
+        self.rate = rate
+        self.curr_calls = 0
+        self.time = 0
+        self.process = None
+        self.stop_event = Event()
+
+    def __repr__(self):
+        return f'Model: {self.name}; Rate: {self.rate}; Current_Calls: {self.curr_calls} calls; Time Passed: {self.time} seconds.'
+    
+    def __str__(self):
+        return self.name
+    
+    def background_task(self):
+        # Background task to manage the rate of calls to the API
+        while not self.stop_event.is_set():
+            time.sleep(5)
+            self.time += 5
+            if self.time >= 60:
+                self.time = 0
+                self.curr_calls = 0
+
+    def start(self):
+        # Start the background task
+        self.process = Process(target=self.background_task)
+        self.process.daemon = True
+        self.process.start()
+        logger.info(f"Background process started with PID: {self.process.pid}")
+
+    def stop(self):
+        # Stop the background task
+        logger.info(f"Stopping background process with PID: {self.process.pid}")
+        self.stop_event.set()
+        if self.process:
+            self.process.join(timeout=5)
+            if self.process.is_alive():
+                self.process.terminate()
+
+class TranslationDataset(Dataset):
+    def __init__(self, texts: List[str], tokenizer, max_length: int = 512):
+        """
+        Custom dataset for efficient text processing
+        
+        Args:
+            texts: List of input texts
+            tokenizer: HuggingFace tokenizer
+            max_length: Maximum sequence length
+        """
+        self.texts = texts
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        return len(self.texts)
+
+    def __getitem__(self, idx):
+        text = self.texts[idx]
+        
+        # Tokenize with padding and truncation
+        encoding = self.tokenizer(
+            text,
+            max_length=self.max_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt'
+        ).to(device)
+        
+        # Remove batch dimension added by tokenizer
+        return {
+            'input_ids': encoding['input_ids'].squeeze(0),
+            'attention_mask': encoding['attention_mask'].squeeze(0)
+        }
+def collate_fn(batch: List[Dict]):
+    """
+    Custom collate function to handle batching
+    """
+    input_ids = torch.stack([item['input_ids'] for item in batch])
+    attention_mask = torch.stack([item['attention_mask'] for item in batch])
+    
+    return {
+        'input_ids': input_ids,
+        'attention_mask': attention_mask
+    }
+
+def generate_text(
+    texts: List[str],
+    model,
+    tokenizer,
+    batch_size: int = 6,  # Smaller batch size uses less VRAM
+    device: str = 'cuda',
+    max_length: int = 512,
+    max_new_tokens: int = 512,
+    **generate_kwargs
+):
+    """
+    Optimized text generation function
+    
+    Args:
+        model: HuggingFace model
+        tokenizer: HuggingFace tokenizer
+        texts: List of input texts
+        batch_size: Batch size for processing
+        device: Device to run inference on
+        max_length: Maximum input sequence length
+        max_new_tokens: Maximum number of tokens to generate
+        generate_kwargs: Additional kwargs for model.generate
+    
+    Returns:
+        List of generated texts
+    """
+    
+    # Create dataset and dataloader
+    dataset = TranslationDataset(texts, tokenizer, max_length)
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        collate_fn=collate_fn
+    )
+    
+    all_generated_texts = []
+    
+    # Default generation parameters
+    generation_config = {
+        'max_new_tokens': max_new_tokens,
+        'num_beams': 4,
+        'do_sample': True,
+        'top_k': 50,
+        'top_p': 0.95,
+        'temperature': 0.7,
+        'no_repeat_ngram_size': 2,
+        'pad_token_id': tokenizer.pad_token_id,
+        'eos_token_id': tokenizer.eos_token_id
+    }
+    
+    # Update with user-provided parameters
+    generation_config.update(generate_kwargs)
+    
+    # Perform generation
+    with torch.no_grad():
+        for batch in dataloader:
+            # Move batch to device
+            batch = {k: v.to(device) for k, v in batch.items()}
+            
+            # Generate text
+            outputs = model.generate(
+                input_ids=batch['input_ids'],
+                attention_mask=batch['attention_mask'],
+                **generation_config
+            )
+            
+            # Decode generated tokens
+            decoded_texts = tokenizer.batch_decode(
+                outputs, 
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=True
+            )
+            
+            all_generated_texts.extend(decoded_texts)
+    
+    return all_generated_texts
+
+if __name__ == '__main__':
+    from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M", local_files_only=True).to(device)
+    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", local_files_only=True)
+    tokenizer.src_lang = "zh"
+    texts = ["你好","我"]
+    print(generate_text(texts,model, tokenizer,  forced_bos_token_id=tokenizer.get_lang_id('en')))

helpers/draw.py

@@ -1,88 +0,0 @@
-from PIL import Image, ImageDraw, ImageFont
-from dotenv import load_dotenv
-import os
-import io
-import numpy as np
-import ast
-from helpers.utils import romanize, intercepts
-load_dotenv()
-
-MAX_TRANSLATE = int(os.getenv('MAX_TRANSLATION', 200))
-FONT_FILE = os.getenv('FONT_FILE')
-FONT_SIZE = int(os.getenv('FONT_SIZE', 16))
-LINE_SPACING = int(os.getenv('LINE_SPACING', 3))
-TEXT_COLOR = os.getenv('TEXT_COLOR', "#ff0000")
-LINE_HEIGHT = FONT_SIZE
-TO_ROMANIZE = ast.literal_eval(os.getenv('TO_ROMANIZE', 'True'))
-font = ImageFont.truetype(FONT_FILE, FONT_SIZE)
-
-
-
-def modify_image_bytes(image_bytes: io.BytesIO, ocr_output, translation: list) -> bytes:
-    # Load the image from bytes
-    with io.BytesIO(image_bytes) as byte_stream:
-        image = Image.open(byte_stream)
-        draw = ImageDraw.Draw(image)
-        translate_image(draw, translation, ocr_output, MAX_TRANSLATE)
-        
-    # Save the modified image back to bytes without changing the format
-    with io.BytesIO() as byte_stream:
-        image.save(byte_stream, format=image.format)  # Save in original format
-        modified_image_bytes = byte_stream.getvalue()
-
-    return modified_image_bytes
-
-def translate_image(draw: ImageDraw, translation: list, ocr_output: list, max_translate: int) -> ImageDraw:
-    translation 
-    translated_number = 0
-    bounding_boxes = []
-    for i, (position, untranslated_phrase, confidence) in enumerate(ocr_output):
-        if translated_number >= max_translate:
-            break
-        translate_one_phrase(draw, translation[i], position, bounding_boxes, untranslated_phrase)
-        translated_number += 1
-    return draw
-
-def translate_one_phrase(draw: ImageDraw, translated_phrase: str, position: tuple, bounding_boxes: list, untranslated_phrase: str) -> ImageDraw:
-    # Draw the bounding box
-    top_left, _, _, _ = position
-    position = (top_left[0], top_left[1] - 60)
-    text_content = f"{translated_phrase}\n{romanize(untranslated_phrase, TO_ROMANIZE)}\n{untranslated_phrase}"
-    lines = text_content.split('\n')
-    x,y = position
-    max_width = 0
-    total_height = 0
-    total_height = len(lines) * (LINE_HEIGHT + LINE_SPACING)
-    for line in lines:
-        bbox = draw.textbbox(position, line, font=font)
-        line_width = bbox[2] - bbox[0]
-        max_width = max(max_width, line_width)
-    bounding_box = (x, y, x + max_width, y + total_height, untranslated_phrase)
-    print(f"Bounding Box of Interest: {bounding_box}")
-
-    adjust_if_intersects(x, y, bounding_box, bounding_boxes, untranslated_phrase, max_width, total_height)
-    adjusted_x, adjusted_y, adjusted_max_x, adjusted_max_y, _ = bounding_boxes[-1]
-    draw.rectangle([(adjusted_x,adjusted_y), (adjusted_max_x, adjusted_max_y)], outline="black", width=1)
-    position = (adjusted_x,adjusted_y)
-    for line in lines:
-        draw.text(position, line, fill= TEXT_COLOR, font=font)
-        adjusted_y += FONT_SIZE + LINE_SPACING
-        position = (adjusted_x,adjusted_y)
-    print(f"Adjusted_bounding_box: {bounding_box[-1]}.\n")
-
-def adjust_if_intersects(x: int, y: int, bounding_box: tuple, bounding_boxes: list, untranslated_phrase: str, max_width: int, total_height: int) -> tuple:
-    y = np.max([y,0])
-    if len(bounding_boxes) > 0:
-        for box in bounding_boxes:
-            print(f'Investigating box: {box}')
-            if intercepts((box[0],box[2]),(bounding_box[0],bounding_box[2])) and intercepts((box[1],box[3]),(y, y+total_height)):
-                print(f'Overlapping change adjustment to {untranslated_phrase}')
-                y = np.max([y,box[3]]) + LINE_SPACING
-                print(y, box[3])
-                print(f'Changed to {(x,y, x+max_width, y+total_height, untranslated_phrase)}')
-    adjusted_bounding_box = (x, y, x + max_width, y + total_height, untranslated_phrase)
-    bounding_boxes.append(adjusted_bounding_box)
-    return adjusted_bounding_box
-
-
-

helpers/ocr.py

@@ -1,20 +1,27 @@
 from paddleocr import PaddleOCR
 import easyocr
+from typing import Optional
 from rapidocr_onnxruntime import RapidOCR
-import langid
-from helpers.utils import contains_lang
+import langid, sys,os
+from utils import contains_lang, standardize_lang
 from concurrent.futures import ThreadPoolExecutor
+sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
+from logging_config import logger
 # PaddleOCR
 # Paddleocr supports Chinese, English, French, German, Korean and Japanese.
 # You can set the parameter `lang` as `ch`, `en`, `fr`, `german`, `korean`, `japan`
 # to switch the language model in order.
 # need to run only once to download and load model into memory
 
-def _paddle_init(lang='ch', use_angle_cls=False, use_GPU=True):
-    return PaddleOCR(use_angle_cls=use_angle_cls, lang=lang, use_GPU=use_GPU)
+default_languages = ['en', 'ch', 'ja', 'ko']
+
+
+def _paddle_init(paddle_lang, use_angle_cls=False, use_GPU=True, **kwargs):
+    return PaddleOCR(use_angle_cls=use_angle_cls, lang=paddle_lang, use_GPU=use_GPU, **kwargs)
 
 
 def _paddle_ocr(ocr, image) -> list:
+
     ### return a list containing the bounding box, text and confidence of the detected text
     result = ocr.ocr(image, cls=False)[0]
     if not isinstance(result, list):
@@ -24,31 +31,34 @@ def _paddle_ocr(ocr, image) -> list:
 
 # EasyOCR has support for many languages
 
-def _easy_init(ocr_languages: list, use_GPU=True):
-    return easyocr.Reader(ocr_languages, gpu=use_GPU)
+def _easy_init(easy_languages: list, use_GPU=True, **kwargs):
+    langs = []
+    for lang in easy_languages:
+        langs.append(standardize_lang(lang)['easyocr_lang'])
+    return easyocr.Reader(langs, gpu=use_GPU, **kwargs)
 
 def _easy_ocr(ocr,image) -> list:
     return ocr.readtext(image)
 
 # RapidOCR mostly for mandarin and some other asian languages
 
-def _rapid_init(use_GPU=True):
-    return RapidOCR(use_gpu=use_GPU)
+def _rapid_init(use_GPU=True, **kwargs):
+    return RapidOCR(use_gpu=use_GPU, **kwargs)
 
 def _rapid_ocr(ocr, image) -> list:
     return ocr(image)
 
 ### Initialize the OCR model
-def init_OCR(model='paddle', **kwargs):
+def init_OCR(model='paddle', easy_languages: Optional[list] = ['ch_sim','en'], paddle_lang: Optional[str] = 'ch', use_GPU=True, **kwargs):
     if model == 'paddle':
-        return _paddle_init(**kwargs)
+        return _paddle_init(paddle_lang=paddle_lang, use_GPU=use_GPU)
     elif model == 'easy':
-        return _easy_init(**kwargs)
+        return _easy_init(easy_languages=easy_languages, use_GPU=use_GPU)
     elif model == 'rapid':
-        return _rapid_init(**kwargs)
+        return _rapid_init(use_GPU=use_GPU)
 
 ### Perform OCR on the image
-def identify(ocr, image) -> list:
+def _identify(ocr, image) -> list:
     if isinstance(ocr, PaddleOCR):
         return _paddle_ocr(ocr, image)
     elif isinstance(ocr, easyocr.Reader):
@@ -56,13 +66,14 @@ def identify(ocr, image) -> list:
     elif isinstance(ocr, RapidOCR):
         return _rapid_ocr(ocr, image)
     else:
-        raise ValueError("Invalid OCR model. Please initialise the OCR model first with init() and pass it as an argument to identify().")
+        raise ValueError("Invalid OCR model. Please initialise the OCR model first with init() and pass it as an argument to _identify().")
     
 
-### Filter out the results that are not in the source language
-def id_filtered(ocr, image, lang) -> list:
-    result = identify(ocr, image)
-
+### Filter out the results that are not in the source language. Slower but for a wider range of languages
+# not working but also not very reliable so don't worry about it
+def _id_filtered(ocr, image, lang) -> list:
+    lang = standardize_lang(lang)['id_model_lang']
+    result = _identify(ocr, image)
     ### Parallelise since langid is slow
     def classify_text(entry):
         return entry if langid.classify(entry[1])[0] == lang else None
@@ -71,12 +82,29 @@ def id_filtered(ocr, image, lang) -> list:
     return results_no_eng
 
 
-# zh, ja, ko
-def id_lang(ocr, image, lang) -> list:
-    result = identify(ocr, image)
-    filtered = [entry for entry in result if contains_lang(entry[1], lang)]
+# ch_sim, ch_tra, ja, ko, en
+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}.")
     return filtered
 
+def id_keep_source_lang(ocr, image, lang) -> list:
+    try:
+        return _id_lang(ocr, image, lang)
+    except ValueError:
+        try:
+            return _id_filtered(ocr, image, lang)
+        except Exception as e:
+            print(f'Probably an issue with the _id_filtered function. {e}')
+            sys.exit(1)
+
+
+
 def get_words(ocr_output) -> list:
     return [entry[1] for entry in ocr_output]
 
@@ -85,3 +113,12 @@ def get_positions(ocr_output) -> list:
 
 def get_confidences(ocr_output) -> list:
     return [entry[2] for entry in ocr_output]
+
+
+if __name__ == '__main__':
+    # OCR_languages = ['ch_sim','en']
+    # image_old = '/home/James/Pictures/Screenshots/DP-1.jpg'
+    # reader = easyocr.Reader(OCR_languages, gpu=True) # this needs to run only once to load the model into memory
+    # result = reader.readtext(image_old)
+    # print(result)
+    print(id_keep_source_lang(init_OCR(model='paddle', paddle_lang='zh', easy_languages=['en', 'ch_sim']), '/home/James/Pictures/Screenshots/DP-1.jpg', 'ch_sim'))

helpers/translation.py

@@ -1,76 +1,217 @@
-from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration,  AutoTokenizer, AutoModelForSeq2SeqLM, GPTQConfig
-import torch, os
-from dotenv import load_dotenv
-load_dotenv()
+from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration,  AutoTokenizer, AutoModelForSeq2SeqLM, GPTQConfig, AutoModelForCausalLM
+import google.generativeai as genai
+import torch, os, sys, ast 
+from utils import standardize_lang
+from functools import wraps
+from batching import generate_text, Gemini
+from logging_config import logger
+from multiprocessing import Process,Event
+# root dir 
+sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
 
-if os.getenv('TRANSLATION_USE_GPU') in ['False', '0', 'false', 'no', 'No', 'NO', 'FALSE']:
-    device = torch.device("cpu")
-else:
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+from config import LOCAL_FILES_ONLY, available_langs, curr_models, BATCH_SIZE, device, GEMINI_KEY, MAX_INPUT_TOKENS, MAX_OUTPUT_TOKENS, seq_llm_models, api_llm_models, causal_llm_models
 
-### Batch translate a list of strings
+##############################
+# translation decorator
+def translate(translation_func):
+    @wraps(translation_func)
+    def wrapper(text, *args, **kwargs):
+        try: 
+            if len(text) == 0:
+                return []
+            return translation_func(text, *args, **kwargs)
+        except Exception as e:
+            logger.error(f"Translation error with the following function: {translation_func.__name__}. Text: {text}\nError: {e}")
+    return wrapper
+###############################
 
 
+###############################
+def init_GEMINI(models_and_rates = None):
+    if not models_and_rates:
+        ## this is default for free tier
+        models_and_rates = {'gemini-1.5-pro': 2, 'gemini-1.5-flash': 15, 'gemini-1.5-flash-8b': 8, 'gemini-1.0-pro': 15} # order from most pref to least pref
+    models = [Gemini(name, rate) for name, rate in models_and_rates.items()]
+    for model in models:
+        model.start()
+    genai.configure(api_key=GEMINI_KEY)
+    return models
+
+def translate_GEMINI(text, models, from_lang, target_lang):
+    safety_settings = {
+    "HARM_CATEGORY_HARASSMENT": "BLOCK_NONE",
+    "HARM_CATEGORY_HATE_SPEECH": "BLOCK_NONE",
+    "HARM_CATEGORY_SEXUALLY_EXPLICIT": "BLOCK_NONE",
+    "HARM_CATEGORY_DANGEROUS_CONTENT": "BLOCK_NONE"}
+    prompt = f"Without any additional remarks, and without any code, translate the following items of the Python list from {from_lang} into {target_lang} and output as a Python list ensuring proper escaping of characters: {text}"
+    for model in models:
+        if model.curr_calls < model.rate:
+            try:
+                response = genai.GenerativeModel(model.name).generate_content(prompt, 
+                                                safety_settings=safety_settings)
+                model.curr_calls += 1   
+                logger.info(repr(model))
+                logger.info(f'Model Response: {response.text.strip()}')
+                return ast.literal_eval(response.text.strip())
+            except Exception as e:
+                logger.error(f"Error with model {model.name}. Error: {e}")
+    logger.error("No models available to translate. Please wait for a model to be available.")
+
+    
+###############################
+# Best model by far. Aya-23-8B. Gemma is relatively good. If I get the time to quantize either gemma or aya those will be good to use. llama3.2 is really good as well.
+def init_AYA():
+    model_id = "CohereForAI/aya-23-8B"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=True)
+    model = AutoModelForCausalLM.from_pretrained(model_id, locals_files_only=True, torch_dtype=torch.float16).to(device)
+    model.eval()
+    return (model, tokenizer)
+
+
+
+##############################
 # M2M100 model
 
 
 def init_M2M():
-    global tokenizer, model
-    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", local_files_only=True)
-    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M", local_files_only=True, torch_dtype=torch.float16).to(device)
+    tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", local_files_only=LOCAL_FILES_ONLY)
+    model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M", local_files_only=LOCAL_FILES_ONLY).to(device)
     model.eval()
+    return (model, tokenizer)
 
 
-
-def translate_M2M(text, from_lang = 'zh', target_lang = 'en'):
+def translate_M2M(text, model, tokenizer, from_lang = 'ch_sim', target_lang = 'en') -> list[str]:
+    model_lang_from = standardize_lang(from_lang)['translation_model_lang']
+    model_lang_to = standardize_lang(target_lang)['translation_model_lang']
     if len(text) == 0:
         return []
-    tokenizer.src_lang = from_lang
-    with torch.no_grad():
-        encoded = tokenizer(text, return_tensors="pt", padding=True).to(device)
-        generated_tokens = model.generate(**encoded,
-                                          forced_bos_token_id=tokenizer.get_lang_id(target_lang))
-    translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
-    return translated
+    tokenizer.src_lang = model_lang_from
+    generated_translations = generate_text(text, model,tokenizer, batch_size=BATCH_SIZE, 
+                  max_length=MAX_INPUT_TOKENS, max_new_tokens=MAX_OUTPUT_TOKENS, 
+                  forced_bos_token_id=tokenizer.get_lang_id(model_lang_to))
+    return generated_translations
+
+###############################
 
 
-# Helsinki-NLP model Opus MT
+###############################
+# Helsinki-NLP model Opus MT 
+# Refer here for all the models https://huggingface.co/Helsinki-NLP
+def get_OPUS_model(from_lang, target_lang):
+    model_lang_from = standardize_lang(from_lang)['translation_model_lang']
+    model_lang_to = standardize_lang(target_lang)['translation_model_lang']
+    return f"Helsinki-NLP/opus-mt-{model_lang_from}-{model_lang_to}"
+        
 
-
-def init_OPUS():
-    global tokenizer, model
-    tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-tc-bible-big-zhx-en", local_files_only=True)
-    model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-tc-bible-big-zhx-en", local_files_only=True, torch_dtype=torch.float16).to(device)
+def init_OPUS(from_lang = 'ch_sim', target_lang = 'en'):
+    opus_model = get_OPUS_model(from_lang, target_lang)
+    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]) -> list[str]:
-    if len(text) == 0:
-        return []
-    with torch.no_grad():
-        encoded = tokenizer(text, return_tensors="pt", padding=True).to(device)
-        generated_tokens = model.generate(**encoded)
-    translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
-    return translated
+def translate_OPUS(text: list[str], model, tokenizer) -> list[str]:
+    translated_text = generate_text(model,tokenizer, text, 
+                  batch_size=BATCH_SIZE, device=device, 
+                  max_length=MAX_INPUT_TOKENS, max_new_tokens=MAX_OUTPUT_TOKENS)
+    return translated_text
 
-### 
-def init_TRANSLATE(model): # model = 'opus' or 'm2m'
-    if model == 'opus':
-        init_OPUS()
-    elif model == 'm2m':
-        init_M2M()
+###############################
+
+
+def init_Seq_LLM(model_type, **kwargs): # model = 'opus' or 'm2m'
+    if model_type == 'opus':
+        return init_OPUS(**kwargs)
+    elif model_type == 'm2m':
+        return init_M2M()
     else:
-        raise ValueError("Invalid model. Please use 'opus' or 'm2m'.")
-    
+        raise ValueError(f"Invalid model. Please use {' or '.join(curr_models)}.")
 
+def init_API_LLM(model_type, **kwargs): # model = 'gemma'
+    if model_type == 'gemini':
+        return init_GEMINI(**kwargs)
+    else:
+        raise ValueError(f"Invalid model type. Please use {' or '.join(api_llm_models)}.")
+
+def init_Causal_LLM(model_type, **kwargs):
+    pass
 ###
-def translate(text, model, **kwargs):
-    if model == 'opus':
-        return translate_OPUS(text)
-    elif model == 'm2m':
+@translate
+def translate_Seq_LLM(text, 
+                       model_type, # 'opus' or 'm2m'
+                        model,
+                        tokenizer,
+                       **kwargs):
+    if model_type == 'opus':
+        return translate_OPUS(text, model, tokenizer)
+    elif model_type == 'm2m':
         try: 
-            return translate_M2M(text, **kwargs)
-        except:
-            raise ValueError("Please provide the from_lang and target_lang variables if you are using the M2M model.")
-            
+            return translate_M2M(text, model, tokenizer, **kwargs)
+        except Exception as e:
+            logger.error(f"Error with M2M model. Error: {e}")
+            # raise ValueError(f"Please provide the correct from_lang and target_lang variables if you are using the M2M model. Use the list from {available_langs}.")
     else:
-        raise ValueError("Invalid model. Please use 'opus' or 'm2m'.")
+        raise ValueError(f"Invalid model. Please use {' or '.join(curr_models)}.")
+
+
+### if you want to use any other translation, just define a translate function with input text and output text.
+# def translate_api(text):
+#@translate 
+#def translate_Causal_LLM(text, model_type, model)
+
+@translate
+def translate_API_LLM(text: list[str], 
+                      model_type: str, # 'gemma'
+                      models: list, # list of objects of classes defined in batching.py
+                       from_lang: str, # suggested to use ISO 639-1 codes
+                       target_lang: str # suggested to use ISO 639-1 codes
+                       ) -> list[str]:
+    if model_type == 'gemini':
+        from_lang = standardize_lang(from_lang)['translation_model_lang']
+        target_lang = standardize_lang(target_lang)['translation_model_lang']
+        return translate_GEMINI(text, models, from_lang, target_lang)
+    else:
+        raise ValueError(f"Invalid model. Please use {' or '.join(api_llm_models)}.")
+
+@translate
+def translate_Causal_LLM(text: list[str],
+                         model_type, # aya
+                            model,
+                            tokenizer,
+                            from_lang: str,
+                            target_lang: str) -> list[str]:
+    model_lang_from = standardize_lang(from_lang)['translation_model_lang']
+    model_lang_to = standardize_lang(target_lang)['translation_model_lang']
+    if len(text) == 0:
+        return []
+    pass
+
+
+# choose between local Seq2Seq LLM or obtain translations from an API
+def init_func(model):
+    if model in seq_llm_models:
+        return init_Seq_LLM
+    elif model in api_llm_models:
+        return init_API_LLM
+    elif model in causal_llm_models:
+        return init_Causal_LLM
+    else:
+        raise ValueError("Invalid model category. Please use either 'seq' or 'api'.")
+
+def translate_func(model):
+    if model in seq_llm_models:
+        return translate_Seq_LLM
+    elif model in api_llm_models:
+        return translate_API_LLM
+    elif model in causal_llm_models:
+        return translate_Causal_LLM
+    else:
+        raise ValueError("Invalid model category. Please use either 'seq' or 'api'.")
+
+
+### todo: if cuda is not detected, default to online translation as cpu just won't cut it bro. Parallel process it over multiple websites to make it faster
+if __name__ == "__main__":
+    models = init_GEMINI()
+    print(translate_API_LLM(['想要借用合成台。', '有什么卖的?', '声音太大会有什么影响吗?', '不怕丢东西吗?', '再见。', '布纳马', '想买什么自己拿。把钱留在旁边就好。', '回顾', '隐藏'], 'gemini', models, from_lang='ch_sim', target_lang='en'))
+    # model, tokenizer = init_M2M()
+    # print(translate_Seq_LLM( ['想要借用合成台。', '有什么卖的?', '声音太大会有什么影响吗?', '不怕丢东西吗?', '再见。', '布纳马', '想买什么自己拿。把钱留在旁边就好。', '回顾', '隐藏'], 'm2m', model, tokenizer, from_lang='ch_sim', target_lang='en'))

helpers/utils.py

@@ -3,12 +3,15 @@ from pypinyin import pinyin
 import pyscreenshot as ImageGrab # wayland tings not sure if it will work on other machines alternatively use mss
 import mss, io, os
 from PIL import Image
-import jaconv
-import MeCab
-import unidic
+import jaconv, MeCab, unidic, pykakasi
+
+# for creating furigana
 mecab = MeCab.Tagger('-d "{}"'.format(unidic.DICDIR))
 uroman = ur.Uroman()
 
+# for romanising japanese text. Can convert to hiragana or katakana as well but does not split the words up so harder to use for furigana
+kks = pykakasi.kakasi()
+
 
 # define a function for checking whether one axis of a shape intercepts with another
 def intercepts(x,y):
@@ -71,7 +74,10 @@ def add_furigana(text):
     furigana_string = '' 
     for i in parsed:
         words = i.split('\t')[0]
-        add = f'({jaconv.kata2hira(i.split(',')[6])})'
+        try :
+            add = f'({jaconv.kata2hira(i.split(',')[6])})'
+        except:
+            add = ''
         to_add = add if contains_kanji(words) else ''
         furigana_string += i.split('\t')[0] + to_add
     return furigana_string
@@ -87,10 +93,12 @@ def contains_katakana(text):
     return bool(re.search(r'[\u30A0-\u30FF]', text))
 
 
-
-def romanize(text, piny=False):
-    if piny:
+# use kakasi to romanize japanese text
+def romanize(text, lang):
+    if lang == 'zh':
         return ' '.join([ py[0] for py in pinyin(text, heteronym=True)])
+    if lang == 'ja':
+        return kks.convert(text)[0]['hepburn']
     return uroman.romanize_string(text)
 
 # check if a string contains words from a language
@@ -107,36 +115,45 @@ def contains_lang(text, lang):
     else:
         raise ValueError("Invalid language. Please use one of 'en', 'zh', 'ja', or 'ko'.")
 
-### en, ch_sim, ja, ko rapidocr only has chinese and en at the moment
+### en, ch_sim, ch_tra, ja, ko rapidocr only has chinese and en at the moment
 def standardize_lang(lang):
     if lang == 'ch_sim':
         easyocr_lang = 'ch_sim'
         paddleocr_lang = 'ch'
         rapidocr_lang = 'ch'
         translation_model_lang = 'zh'
+        id_model_lang = 'zh'
     elif lang == 'ch_tra':
         easyocr_lang = 'ch_tra'
         paddleocr_lang = 'ch'
         rapidocr_lang = 'ch'
         translation_model_lang = 'zh'
+        id_model_lang = 'zh'
     elif lang == 'ja':
         easyocr_lang = 'ja'
         paddleocr_lang = 'ja'
         rapidocr_lang = 'ja'
         translation_model_lang = 'ja'
+        id_model_lang = 'ja'
     elif lang == 'ko':
         easyocr_lang = 'korean'
         paddleocr_lang = 'ko'
         rapidocr_lang = 'ko'
         translation_model_lang = 'ko'
+        id_model_lang = 'ko'
     elif lang == 'en':
         easyocr_lang = 'en'
         paddleocr_lang = 'en'
         rapidocr_lang = 'en'
         translation_model_lang = 'en'
+        id_model_lang = 'en'
     else:
-        raise ValueError("Invalid language. Please use one of 'en', 'ch_sim', 'ch_tra', 'ja', or 'ko'.")
-    return {'easyocr_lang': easyocr_lang, 'paddleocr_lang': paddleocr_lang, 'rapidocr_lang': rapidocr_lang, 'translation_model_lang': translation_model_lang}
+        raise ValueError(f"Invalid language {lang}. Please use one of 'en', 'ch_sim', 'ch_tra', 'ja', or 'ko'.")
+    return {'easyocr_lang': easyocr_lang, 
+            'paddleocr_lang': paddleocr_lang, 
+            'rapidocr_lang': rapidocr_lang, 
+            'translation_model_lang': translation_model_lang,
+            'id_model_lang': id_model_lang}
 
 def which_ocr_lang(model):
     if model == 'easy':