chickenflyshigh/onscreen-translator
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Batched asynchronous API requests. Added additional draw options.

Browse Source
This commit is contained in:
chickenflyshigh
2024-11-06 22:51:03 +11:00
parent 499a2c3972
commit 11600ae70f
17 changed files with 788 additions and 711 deletions
Download Patch File Download Diff File Expand all files Collapse all files
helpers/batching.py
+284 -100
View File
@@ -1,142 +1,326 @@
from torch.utils.data import Dataset, DataLoader
from typing import List, Dict
from datetime import datetime, timedelta
from dotenv import load_dotenv
import os , sys, torch, time, ast
import os , sys, torch, time, ast, json, pytz
from werkzeug.exceptions import TooManyRequests
from multiprocessing import Process, Event, Value
load_dotenv()
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from config import device, GEMINI_API_KEY, GROQ_API_KEY
from config import device, GEMINI_API_KEY, GROQ_API_KEY, MAX_TRANSLATE
from logging_config import logger
from groq import Groq
from groq import Groq as Groqq
import google.generativeai as genai
from google.api_core.exceptions import ResourceExhausted
import asyncio
import aiohttp
from functools import wraps
from data import session, Api, Translations
from typing import Optional
class ApiModel():
def __init__(self, model, # model name
rate, # rate of calls per minute
api_key, # api key for the model wrt the site
def __init__(self, model, # model name as defined by the API
site, # site of the model; # to be precise, use the name as defined precisely by the class names in this script, i.e. Groqq and Gemini
api_key: Optional[str] = None, # api key for the model wrt the site
rpmin: Optional[int] = None, # rate of calls per minute
rph: Optional[int] = None, # rate of calls per hour
rpd: Optional[int] = None, # rate of calls per day
rpw: Optional[int] = None, # rate of calls per week
rpmth: Optional[int] = None, # rate of calls per month
rpy: Optional[int] = None # rate of calls per year
):
self.model = model
self.rate = rate
self.api_key = api_key
self.curr_calls = Value('i', 0)
self.time = Value('i', 0)
self.process = None
self.stop_event = Event()
self.site = None
self.model = model
self.rpmin = rpmin
self.rph = rph
self.rpd = rpd
self.rpw = rpw
self.rpmth = rpmth
self.rpy = rpy
self.site = site
self.from_lang = None
self.target_lang = None
self.request = None # request response from API
self.db_table = None
self.session_calls = 0
self._id = None
self._set_db_model_id() if self._get_db_model_id() else self.update_db()
# Create the table if it does not already exist
def __repr__(self):
return f'{self.site} Model: {self.model}; Rate: {self.rate}; Current_Calls: {self.curr_calls.value} calls; Time Passed: {self.time.value} seconds.'
return f'{self.site} Model: {self.model}; Total calls this session: {self.session_calls}; rpmin: {self.rpmin}; rph: {self.rph}; rpd: {self.rpd}; rpw: {self.rpw}; rpmth: {self.rpmth}; rpy: {self.rpy}'
def __str__(self):
return self.model
async def __aenter__(self):
self.session = aiohttp.ClientSession()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
if self.session:
await self.session.close()
def _get_db_model_id(self):
model = session.query(Api).filter_by(model_name = self.model, site = self.site).first()
if model:
return model.id
else:
return None
def _set_db_model_id(self):
self._id = self._get_db_model_id()
@staticmethod
def _get_time():
return datetime.now(tz=pytz.timezone('Australia/Sydney'))
def set_lang(self, from_lang, target_lang):
self.from_lang = from_lang
self.target_lang = target_lang
### CHECK MINUTELY API RATES. For working with hourly rates and monthly will need to create another file. Also just unlikely those rates will be hit
async def api_rate_check(self):
# Background task to manage the rate of calls to the API
while not self.stop_event.is_set():
start_time = time.monotonic()
self.time.value += 5
if self.time.value >= 60:
self.time.value = 0
self.curr_calls.value = 0
elapsed = time.monotonic() - start_time
# Sleep for exactly 5 seconds minus the elapsed time
sleep_time = max(0, 5 - elapsed)
await asyncio.sleep(sleep_time)
def set_db_table(self, db_table):
self.db_table = db_table
def update_db(self):
api = session.query(Api).filter_by(model_name = self.model, site = self.site).first()
if not api:
api = Api(model_name = self.model,
site = self.site,
rpmin = self.rpmin,
rph = self.rph,
rpd = self.rpd,
rpw = self.rpw,
rpmth = self.rpmth,
rpy = self.rpy)
session.add(api)
session.commit()
self._set_db_model_id()
else:
api.rpmin = self.rpmin
api.rph = self.rph
api.rpd = self.rpd
api.rpw = self.rpw
api.rpmth = self.rpmth
api.rpy = self.rpy
session.commit()
def _db_add_translation(self, text: list | str, translation: list, mismatch = False):
text = json.dumps(text) if isinstance(text, list) else json.dumps([text])
translation = json.dumps(translation)
translation = Translations(source_texts = text, translated_texts = translation,
model_id = self._id, source_lang = self.from_lang, target_lang = self.target_lang,
timestamp = datetime.now(tz=pytz.timezone('Australia/Sydney')),
translation_mismatch = mismatch)
session.add(translation)
session.commit()
@staticmethod
def _single_period_calls_check(max_calls, call_count):
if not max_calls:
return True
if max_calls <= call_count:
return False
else:
return True
def _are_rates_good(self):
curr_time = self._get_time()
min_ago = curr_time - timedelta(minutes=1)
hour_ago = curr_time - timedelta(hours=1)
day_ago = curr_time - timedelta(days=1)
week_ago = curr_time - timedelta(weeks=1)
month_ago = curr_time - timedelta(days=30)
year_ago = curr_time - timedelta(days=365)
min_calls = session.query(Translations).join(Api). \
filter(Api.id==self._id,
Translations.timestamp >= min_ago
).count()
hour_calls = session.query(Translations).join(Api). \
filter(Api.id==self._id,
Translations.timestamp >= hour_ago
).count()
day_calls = session.query(Translations).join(Api). \
filter(Api.id==self._id,
Translations.timestamp >= day_ago
).count()
week_calls = session.query(Translations).join(Api). \
filter(Api.id==self._id,
Translations.timestamp >= week_ago
).count()
month_calls = session.query(Translations).join(Api). \
filter(Api.id==self._id,
Translations.timestamp >= month_ago
).count()
year_calls = session.query(Translations).join(Api). \
filter(Api.id==self._id,
Translations.timestamp >= year_ago
).count()
if self._single_period_calls_check(self.rpmin, min_calls) \
and self._single_period_calls_check(self.rph, hour_calls) \
and self._single_period_calls_check(self.rpd, day_calls) \
and self._single_period_calls_check(self.rpw, week_calls) \
and self._single_period_calls_check(self.rpmth, month_calls) \
and self._single_period_calls_check(self.rpy, year_calls):
return True
else:
logger.warning(f"Rate limit reached for {self.model} from {self.site}. Current calls: {min_calls} in the last minute; {hour_calls} in the last hour; {day_calls} in the last day; {week_calls} in the last week; {month_calls} in the last month; {year_calls} in the last year.")
return False
# async def request_func(request):
# @wraps(request)
# async def wrapper(self, text, *args, **kwargs):
# if await self._are_rates_good():
# try:
# self.session_calls += 1
# response = await request(self, text, *args, **kwargs)
# return response
# except Exception as e:
# logger.error(f"Error with model {self.model} from {self.site}. Error: {e}")
# else:
# logger.error(f"Rate limit reached for this model.")
# raise TooManyRequests('Rate limit reached for this model.')
# return wrapper
def background_task(self):
asyncio.run(self.api_rate_check())
# @request_func
async def translate(self, texts_to_translate, store = False):
if isinstance(texts_to_translate, str):
texts_to_translate = [texts_to_translate]
if len(texts_to_translate) == 0:
return []
#prompt = f"Without any additional remarks, and without any code, translate the following items of the Python list from {self.from_lang} into {self.target_lang} and output as a Python list ensuring proper escaping of characters and ensuring the length of the list given is exactly equal to the length of the list you provide. Do not output in any other language other than the specified target language: {texts_to_translate}"
prompt = f"""INSTRUCTIONS:
- Provide ONE and ONLY ONE translation to each text provided in the JSON array given.
- The translations must preserve the original order.
- Each translation must be from the Source language to the Target language
- Source language: {self.from_lang}
- Target language: {self.target_lang}
- Texts are provided in JSON array syntax.
- Respond using ONLY valid JSON array syntax.
- Do not include explanations or additional text
- Escape special characters properly
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}")
Input texts:
{texts_to_translate}
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()
Expected format:
["translation1", "translation2", ...]
def request_func(request):
@wraps(request)
def wrapper(self, text, *args, **kwargs):
if self.curr_calls.value < self.rate:
# try:
response = request(self, text, *args, **kwargs)
self.curr_calls.value += 1
return response
# except Exception as e:
#logger.error(f"Error with model {self.model} from {self.site}. Error: {e}")
else:
logger.error(f"Rate limit reached for this model. Please wait for the rate to reset in {60 - self.time} seconds.")
raise TooManyRequests('Rate limit reached for this model.')
return wrapper
Translation:"""
response = await self._request(prompt)
response_list = ast.literal_eval(response.strip())
logger.debug(repr(self))
logger.info(f'{self.model} translated texts from: {texts_to_translate} to {response_list}.')
if len(response_list) != len(texts_to_translate) and len(texts_to_translate) <= MAX_TRANSLATE:
logger.error(f"{self.model} model failed to translate all the texts. Number of translations to make: {len(texts_to_translate)}; Number of translated texts: {len(response_list)}.")
if store:
self._db_add_translation(texts_to_translate, response_list, mismatch=True)
else:
if store:
self._db_add_translation(texts_to_translate, response_list)
print(response_list)
return response_list
class Groq(ApiModel):
def __init__(self, # model name as defined by the API
model,
api_key = GROQ_API_KEY, # api key for the model wrt the site
**kwargs):
super().__init__(model,
api_key = api_key,
site = 'Groq', **kwargs)
self.client = Groqq()
@request_func
def translate(self, request_fn, texts_to_translate):
if len(texts_to_translate) == 0:
return []
prompt = f"Without any additional remarks, and without any code, translate the following items of the Python list from {self.from_lang} into {self.target_lang} and output as a Python list ensuring proper escaping of characters: {texts_to_translate}"
response = request_fn(self, prompt)
return ast.literal_eval(response.strip())
class Groqq(ApiModel):
def __init__(self, model, rate, api_key = GROQ_API_KEY):
super().__init__(model, rate, api_key)
self.site = "Groq"
def request(self, content):
client = Groq()
chat_completion = client.chat.completions.create(
messages=[
{
"role": "user",
"content": content,
async def _request(self, content: str) -> str:
async with aiohttp.ClientSession() as session:
async with session.post(
"https://api.groq.com/openai/v1/chat/completions",
headers={
"Authorization": f"Bearer {GROQ_API_KEY}",
"Content-Type": "application/json"
},
json={
"messages": [{"role": "user", "content": content}],
"model": self.model
}
],
model=self.model
)
return chat_completion.choices[0].message.content
) as response:
response_json = await response.json()
return response_json["choices"][0]["message"]["content"]
# https://console.groq.com/settings/limits for limits
# def request(self, content):
# chat_completion = self.client.chat.completions.create(
# messages=[
# {
# "role": "user",
# "content": content,
# }
# ],
# model=self.model
# )
# return chat_completion.choices[0].message.content
def translate(self, texts_to_translate):
return super().translate(Groqq.request, texts_to_translate)
# async def translate(self, texts_to_translate):
# return super().translate(self.request, texts_to_translate)
class Gemini(ApiModel):
def __init__(self, model, rate, api_key = GEMINI_API_KEY):
super().__init__(model, rate, api_key)
self.site = "Gemini"
def __init__(self, # model name as defined by the API
model,
api_key = GEMINI_API_KEY, # api key for the model wrt the site
**kwargs):
super().__init__(model,
api_key = api_key,
site = 'Google',
**kwargs)
def request(self, content):
genai.configure(api_key=self.api_key)
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"}
response = genai.GenerativeModel(self.model).generate_content(content, safety_settings=safety_settings)
return response.text.strip()
# def request(self, content):
# genai.configure(api_key=self.api_key)
# 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"}
# try:
# response = genai.GenerativeModel(self.model).generate_content(content, safety_settings=safety_settings)
# except ResourceExhausted as e:
# logger.error(f"Rate limited with {self.model}. Error: {e}")
# raise ResourceExhausted("Rate limited.")
# return response.text.strip()
def translate(self, texts_to_translate):
return super().translate(Gemini.request, texts_to_translate)
async def _request(self, content):
async with aiohttp.ClientSession() as session:
async with session.post(
f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={self.api_key}",
headers={
"Content-Type": "application/json"
},
json={
"contents": [{"parts": [{"text": content}]}],
"safetySettings": [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE",
"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE",
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE",
"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"
}
]
}
) as response:
response_json = await response.json()
return response_json['candidates'][0]['content']['parts'][0]['text']
# async def translate(self, texts_to_translate):
# return super().translate(self.request, texts_to_translate)
###################################################################################################
### LOCAL LLM TRANSLATION
class TranslationDataset(Dataset):
def __init__(self, texts: List[str], tokenizer, max_length: int = 512):
"""
helpers/ocr.py
+11 -10
View File
@@ -21,7 +21,6 @@ def _paddle_init(paddle_lang, use_angle_cls=False, use_GPU=True, **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):
@@ -32,28 +31,29 @@ def _paddle_ocr(ocr, image) -> list:
# EasyOCR has support for many languages
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)
return easyocr.Reader(easy_languages, gpu=use_GPU, **kwargs)
def _easy_ocr(ocr,image) -> list:
return ocr.readtext(image)
# RapidOCR mostly for mandarin and some other asian languages
# default only supports chinese and english
def _rapid_init(use_GPU=True, **kwargs):
return RapidOCR(use_gpu=use_GPU, **kwargs)
def _rapid_ocr(ocr, image) -> list:
return ocr(image)
return ocr(image)[0]
### Initialize the OCR model
def init_OCR(model='paddle', easy_languages: Optional[list] = ['ch_sim','en'], paddle_lang: Optional[str] = 'ch', use_GPU=True, **kwargs):
def init_OCR(model='paddle', easy_languages: Optional[list] = ['ch_sim','en'], paddle_lang: Optional[str] = 'ch_sim', use_GPU=True):
if model == 'paddle':
paddle_lang = standardize_lang(paddle_lang)['paddleocr_lang']
return _paddle_init(paddle_lang=paddle_lang, use_GPU=use_GPU)
elif model == 'easy':
return _easy_init(easy_languages=easy_languages, use_GPU=use_GPU)
langs = []
for lang in easy_languages:
langs.append(standardize_lang(lang)['easyocr_lang'])
return _easy_init(easy_languages=langs, use_GPU=use_GPU)
elif model == 'rapid':
return _rapid_init(use_GPU=use_GPU)
@@ -82,10 +82,11 @@ def _id_filtered(ocr, image, lang) -> list:
return results_no_eng
# ch_sim, ch_tra, ja, ko, en
# ch_sim, ch_tra, ja, ko, en input
def _id_lang(ocr, image, lang) -> list:
result = _identify(ocr, image)
lang = standardize_lang(lang)['id_model_lang']
print(result)
try:
filtered = [entry for entry in result if contains_lang(entry[1], lang)]
except:
helpers/translation.py
+54 -16
View File
@@ -1,17 +1,16 @@
from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, GPTQConfig, AutoModelForCausalLM
import google.generativeai as genai
import torch, os, sys, ast, json
import torch, os, sys, ast, json, asyncio, batching, random
from typing import List, Optional, Set
from utils import standardize_lang
from functools import wraps
import random
import batching
from batching import generate_text, Gemini, Groq
from batching import generate_text, Gemini, Groq, ApiModel
from logging_config import logger
from multiprocessing import Process,Event
from asyncio import Task
# root dir
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from config import LOCAL_FILES_ONLY, available_langs, curr_models, BATCH_SIZE, device, GEMINI_API_KEY, MAX_INPUT_TOKENS, MAX_OUTPUT_TOKENS, seq_llm_models, api_llm_models, causal_llm_models
from config import LOCAL_FILES_ONLY, available_langs, curr_models, BATCH_SIZE, device, GEMINI_API_KEY, MAX_INPUT_TOKENS, MAX_OUTPUT_TOKENS, seq_llm_models, api_llm_models, causal_llm_models, API_MODELS_FILEPATH
##############################
# translation decorator
@@ -32,27 +31,66 @@ def translate(translation_func):
def init_API_LLM(from_lang, target_lang):
from_lang = standardize_lang(from_lang)['translation_model_lang']
target_lang = standardize_lang(target_lang)['translation_model_lang']
with open('api_models.json', 'r') as f:
with open(API_MODELS_FILEPATH, 'r') as f:
models_and_rates = json.load(f)
models = []
for class_type, class_models in models_and_rates.items():
cls = getattr(batching, class_type)
instantiated_objects = [ cls(model, rate) for model, rate in class_models.items()]
instantiated_objects = [ cls(model = model, **rates) for model, rates in class_models.items()]
models.extend(instantiated_objects)
for model in models:
model.start()
model.update_db()
model.set_lang(from_lang, target_lang)
return models
def translate_API_LLM(text, models):
random.shuffle(models)
for model in models:
async def translate_API_LLM(texts_to_translate: List[str],
models: List[ApiModel],
call_size: int = 2,
stagger_delay: int = 2) -> List[str]:
async def try_translate(model: ApiModel) -> Optional[List[str]]:
try:
return model.translate(text)
except:
continue
result = await model.translate(texts_to_translate, store=True)
logger.debug(f'Try_translate result: {result}')
return result
except Exception as e:
logger.error(f"Translation failed for {model.model} from {model.site}: {e}")
return None
random.shuffle(models)
groups = [models[i:i+call_size] for i in range(0, len(models), call_size)]
for group in groups:
tasks = set(asyncio.create_task(try_translate(model)) for model in group)
while tasks:
done, pending = await asyncio.wait(tasks,
return_when=asyncio.FIRST_COMPLETED
)
logger.debug(f"Tasks done: {done}")
logger.debug(f"Tasks remaining: {pending}")
for task in done:
result = await task
logger.debug(f'Result: {result}')
if result is not None:
# Cancel remaining tasks
for t in pending:
t.cancel()
return result
logger.error("All models have failed to translate the text.")
raise TypeError("Models have likely all outputted garbage translations or rate limited.")
# def translate_API_LLM(text, models):
# random.shuffle(models)
# logger.debug(f"All Models Available: {models}")
# for model in models:
# logger.info(f"Attempting translation with model {model}.")
# try:
# translation = model.translate(text)
# logger.debug(f"Translation obtained: {translation}")
# if translation or translation == []:
# return translation
# except Exception as e:
# logger.error(f"Error with model {repr(model)}. Error: {e}")
# continue
# logger.error("All models have failed to translate the text.")
# raise TypeError("Models have likely all outputted garbage translations or rate limited.")
###############################
# 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.
helpers/utils.py
+23 -4
View File
@@ -4,7 +4,9 @@ import pyscreenshot as ImageGrab # wayland tings not sure if it will work on oth
import mss, io, os
from PIL import Image
import jaconv, MeCab, unidic, pykakasi
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
import numpy as np
# for creating furigana
mecab = MeCab.Tagger('-d "{}"'.format(unidic.DICDIR))
uroman = ur.Uroman()
@@ -95,7 +97,7 @@ def contains_katakana(text):
# use kakasi to romanize japanese text
def romanize(text, lang):
if lang == 'zh':
if lang in ['zh','ch_sim','ch_tra']:
return ' '.join([ py[0] for py in pinyin(text, heteronym=True)])
if lang == 'ja':
return kks.convert(text)[0]['hepburn']
@@ -131,13 +133,13 @@ def standardize_lang(lang):
id_model_lang = 'zh'
elif lang == 'ja':
easyocr_lang = 'ja'
paddleocr_lang = 'ja'
paddleocr_lang = 'japan'
rapidocr_lang = 'ja'
translation_model_lang = 'ja'
id_model_lang = 'ja'
elif lang == 'ko':
easyocr_lang = 'korean'
paddleocr_lang = 'ko'
paddleocr_lang = 'korean'
rapidocr_lang = 'ko'
translation_model_lang = 'ko'
id_model_lang = 'ko'
@@ -165,6 +167,23 @@ def which_ocr_lang(model):
else:
raise ValueError("Invalid OCR model. Please use one of 'easy', 'paddle', or 'rapid'.")
def similar_tfidf(list1,list2,threshold) -> float:
"""Calculate cosine similarity using TF-IDF vectors."""
if not list1 or not list2:
return 0.0
vectorizer = TfidfVectorizer()
all_texts = list1 + list2
tfidf_matrix = vectorizer.fit_transform(all_texts)
# Calculate average vectors for each list
vec1 = np.mean(tfidf_matrix[:len(list1)].toarray(), axis=0).reshape(1, -1)
vec2 = np.mean(tfidf_matrix[len(list1):].toarray(), axis=0).reshape(1, -1)
return float(cosine_similarity(vec1, vec2)[0, 0]) > threshold
if __name__ == "__main__":
# Example usage