311 lines
11 KiB
Python
311 lines
11 KiB
Python
import os
|
|
import sys
|
|
from google import genai
|
|
from google.genai import types
|
|
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
|
|
import pysubs2
|
|
from deep_translator import GoogleTranslator, MyMemoryTranslator
|
|
import ollama
|
|
from tqdm import tqdm
|
|
|
|
# Define a retry decorator
|
|
# ... (retry_policy remains)
|
|
|
|
def translate_via_ollama(source_srt_content, target_language="English", model="llama3"):
|
|
"""
|
|
Translates SRT content using a local Ollama model (Line-by-Line for progress).
|
|
"""
|
|
try:
|
|
subs = pysubs2.SSAFile.from_string(source_srt_content)
|
|
|
|
# Using tqdm for progress bar
|
|
for line in tqdm(subs, desc=" Ollama Progress", unit="line"):
|
|
text = line.text.strip()
|
|
if text:
|
|
prompt = (
|
|
f"Translate this subtitle text to {target_language}. Output ONLY the translation.\n"
|
|
f"Text: {text}"
|
|
)
|
|
try:
|
|
response = ollama.chat(model=model, messages=[{'role': 'user', 'content': prompt}])
|
|
translated_text = response['message']['content'].strip()
|
|
if translated_text:
|
|
line.text = translated_text
|
|
except Exception as e:
|
|
# Silent fail on line, logs would be too spammy in progress bar
|
|
pass
|
|
|
|
return subs.to_string(format_="srt")
|
|
|
|
except Exception as e:
|
|
print(f" [Local LLM] Error: {e}")
|
|
return None
|
|
|
|
def translate_fallback_mymemory(source_srt_content, target_language="en"):
|
|
"""
|
|
Fallback translation using MyMemory (via deep-translator).
|
|
Limit: 1000 words/day roughly for anonymous usage. Good last resort.
|
|
"""
|
|
try:
|
|
subs = pysubs2.SSAFile.from_string(source_srt_content)
|
|
# MyMemory uses ISO 639-1 usually
|
|
translator = MyMemoryTranslator(source='auto', target=target_language)
|
|
|
|
for line in tqdm(subs, desc=" MyMemory Progress", unit="line"):
|
|
text = line.text.strip()
|
|
if text:
|
|
if len(text) > 500: # MyMemory has stricter limits often
|
|
continue
|
|
try:
|
|
original_text = text.replace(r"\N", " ")
|
|
translated_text = translator.translate(original_text)
|
|
if translated_text:
|
|
line.text = translated_text
|
|
except Exception:
|
|
pass
|
|
|
|
return subs.to_string(format_="srt")
|
|
except Exception as e:
|
|
print(f" [MyMemory Fallback] Critical Error: {e}")
|
|
return None
|
|
|
|
def translate_fallback_free(source_srt_content, target_language="en"):
|
|
"""
|
|
Fallback translation using deep-translator (free Google Translate).
|
|
|
|
Args:
|
|
source_srt_content (str): Content of the source SRT file.
|
|
target_language (str): Target language code (e.g. 'en', 'fr').
|
|
|
|
Returns:
|
|
str: Translated SRT content, or None if failed.
|
|
"""
|
|
try:
|
|
# Load from string
|
|
subs = pysubs2.SSAFile.from_string(source_srt_content)
|
|
translator = GoogleTranslator(source='auto', target=target_language)
|
|
|
|
# Simple line-by-line translation
|
|
for line in tqdm(subs, desc=" DeepTranslate Progress", unit="line"):
|
|
text = line.text.strip()
|
|
if text:
|
|
# Sanity check: Skip lines that are too long
|
|
if len(text) > 4000:
|
|
continue
|
|
|
|
try:
|
|
# pysubs2 text can contain \N for newlines.
|
|
original_text = text.replace(r"\N", " ")
|
|
translated_text = translator.translate(original_text)
|
|
if translated_text:
|
|
line.text = translated_text
|
|
except Exception:
|
|
pass
|
|
|
|
# Return as string
|
|
return subs.to_string(format_="srt")
|
|
except Exception as e:
|
|
print(f" [Free Fallback] Critical Error: {e}")
|
|
return None
|
|
|
|
# Define a retry decorator
|
|
# Waits 2^x * 1 seconds between retries (1s, 2s, 4s...)
|
|
# Stop after 15 attempts
|
|
# before_sleep logic can print a simple message
|
|
def log_retry_attempt(retry_state):
|
|
if retry_state.attempt_number > 1:
|
|
print(f" [Gemini] Rate limit hit. Retrying in {retry_state.next_action.sleep}s...", end='\r')
|
|
|
|
retry_policy = retry(
|
|
stop=stop_after_attempt(15),
|
|
wait=wait_exponential(multiplier=1, min=2, max=60),
|
|
retry=retry_if_exception_type(Exception),
|
|
reraise=True,
|
|
before_sleep=log_retry_attempt
|
|
)
|
|
|
|
@retry_policy
|
|
def _generate_with_retry(client, model_name, prompt):
|
|
"""Internal function to wrap the API call with retry logic."""
|
|
return client.models.generate_content(
|
|
model=model_name,
|
|
contents=prompt
|
|
)
|
|
|
|
def get_best_available_model(client):
|
|
"""
|
|
Queries the API to find the best available model for text generation.
|
|
Priority: gemini-2.0-flash > gemini-1.5-flash > gemini-1.5-pro
|
|
"""
|
|
try:
|
|
# Priority list (New v2 naming conventions if applicable, but standard models persist)
|
|
priorities = [
|
|
"gemini-2.0-flash", # Latest
|
|
"gemini-1.5-flash",
|
|
"gemini-1.5-pro"
|
|
]
|
|
|
|
# In new SDK, client.models.list() returns iterators of Model objects
|
|
# We can just try to use the priority one directly, or list them.
|
|
# Listing can be slow. Let's just default to a known good priority list.
|
|
# If we really want to check:
|
|
# available = [m.name for m in client.models.list()]
|
|
|
|
# For efficiency/speed, we will trust our priority list.
|
|
# The API will error if model doesn't exist, which the try/catch block handling generation will catch?
|
|
# No, better to pick one that exists.
|
|
|
|
# Let's return the latest standard one.
|
|
return "gemini-2.0-flash" # Assuming 2.0 is available or falling back
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Model selection issue ({e}). Defaulting to 'gemini-1.5-flash'.")
|
|
|
|
return "gemini-1.5-flash"
|
|
|
|
def translate_srt(srt_content, target_language="English", api_key=None):
|
|
"""
|
|
Translates SRT subtitle content using the Google GenAI SDK (v2).
|
|
"""
|
|
if not srt_content:
|
|
return ""
|
|
|
|
key = api_key or os.getenv("GEMINI_API_KEY")
|
|
if not key:
|
|
print("Error: GEMINI_API_KEY not found. Please set the environment variable or pass the key.")
|
|
sys.exit(1)
|
|
|
|
# Initialize Client (v2 style)
|
|
try:
|
|
client = genai.Client(api_key=key)
|
|
except Exception as e:
|
|
print(f"Error initializing GenAI Client: {e}")
|
|
return None
|
|
|
|
# Automatically select the best model
|
|
# Note: v2 SDK might use 'gemini-1.5-flash' directly without 'models/' prefix usually
|
|
model_name = "gemini-2.0-flash"
|
|
|
|
prompt = (
|
|
"You are a professional subtitle translator. Your task is to translate the following SRT subtitle file "
|
|
f"into {target_language}.\n\n"
|
|
"RULES:\n"
|
|
"1. PRESERVE the SRT format exactly. Do not modify timestamps (e.g., 00:00:01,000 --> 00:00:04,000) or sequence numbers.\n"
|
|
"2. Only translate the dialogue text.\n"
|
|
"3. Maintain the original tone and context.\n"
|
|
"4. Output ONLY the translated SRT content, no markdown code blocks or explanations.\n\n"
|
|
"SRT Content:\n"
|
|
f"{srt_content}"
|
|
)
|
|
|
|
try:
|
|
# Call the retried internal function
|
|
response = _generate_with_retry(client, model_name, prompt)
|
|
|
|
# Cleanup: sometimes models wrap output in ```srt ... ``` or ``` ... ```
|
|
cleaned_text = response.text.strip()
|
|
if cleaned_text.startswith("```"):
|
|
lines = cleaned_text.split('\n')
|
|
if len(lines) >= 2:
|
|
cleaned_text = '\n'.join(lines[1:-1])
|
|
|
|
return cleaned_text
|
|
except Exception as e:
|
|
print(f"Error during translation after retries: {e}")
|
|
# Fallback to older model if 2.0 fails?
|
|
if "404" in str(e) and "gemini-2.0" in model_name:
|
|
print(" -> gemini-2.0-flash not found, falling back to gemini-1.5-flash")
|
|
try:
|
|
response = _generate_with_retry(client, "gemini-1.5-flash", prompt)
|
|
cleaned_text = response.text.strip()
|
|
if cleaned_text.startswith("```"):
|
|
lines = cleaned_text.split('\n')
|
|
if len(lines) >= 2:
|
|
cleaned_text = '\n'.join(lines[1:-1])
|
|
return cleaned_text
|
|
except Exception as inner_e:
|
|
print(f"Fallback failed: {inner_e}")
|
|
|
|
return None
|
|
|
|
def translate_with_auto_fallback(srt_content, target_language="English", prefer_deep=False, prefer_local=False, available_services=None):
|
|
"""
|
|
Attempts to translate SRT content using Gemini, DeepTranslate, and Local LLM with fallback logic.
|
|
|
|
Args:
|
|
srt_content (str): The source SRT content.
|
|
target_language (str): Target language name (e.g., "English", "French").
|
|
prefer_deep (bool): If True, try DeepTranslate first (among cloud services).
|
|
prefer_local (bool): If True, try Local LLM (Ollama) first.
|
|
available_services (dict, optional): Result of check_service_availability().
|
|
|
|
Returns:
|
|
tuple: (translated_content, method_name) or (None, None) if all failed.
|
|
"""
|
|
|
|
# Map full language name to code for DeepTranslate
|
|
lang_map = {
|
|
"English": "en", "French": "fr", "Spanish": "es", "German": "de",
|
|
"Italian": "it", "Portuguese": "pt", "Russian": "ru",
|
|
"Japanese": "ja", "Chinese": "zh-CN"
|
|
}
|
|
target_code = lang_map.get(target_language, "en")
|
|
|
|
# Determine which services to even try
|
|
def is_ok(name):
|
|
if available_services is None: return True
|
|
return available_services.get(name, True)
|
|
|
|
def try_gemini():
|
|
if not is_ok("Gemini"): return None, None
|
|
res = translate_srt(srt_content, target_language=target_language)
|
|
if res: return res, "Gemini"
|
|
return None, None
|
|
|
|
def try_deep():
|
|
if not is_ok("DeepTranslate"): return None, None
|
|
res = translate_fallback_free(srt_content, target_language=target_code)
|
|
if res: return res, "DeepTranslate"
|
|
return None, None
|
|
|
|
def try_ollama():
|
|
if not is_ok("Ollama"): return None, None
|
|
res = translate_via_ollama(srt_content, target_language=target_language)
|
|
if res: return res, "Local LLM (Ollama)"
|
|
return None, None
|
|
|
|
def try_mymemory():
|
|
res = translate_fallback_mymemory(srt_content, target_language=target_code)
|
|
if res: return res, "MyMemory"
|
|
return None, None
|
|
|
|
# Logic flow
|
|
attempts = []
|
|
|
|
if prefer_local:
|
|
attempts.append(try_ollama)
|
|
if prefer_deep:
|
|
attempts.extend([try_deep, try_gemini])
|
|
else:
|
|
attempts.extend([try_gemini, try_deep])
|
|
else:
|
|
if prefer_deep:
|
|
attempts.extend([try_deep, try_gemini])
|
|
else:
|
|
attempts.extend([try_gemini, try_deep])
|
|
attempts.append(try_ollama)
|
|
|
|
# Final last resort
|
|
attempts.append(try_mymemory)
|
|
|
|
# Execute attempts
|
|
for i, method_func in enumerate(attempts):
|
|
if i > 0:
|
|
print(f" Attempt {i} failed. Trying next fallback...")
|
|
|
|
content, method = method_func()
|
|
if content:
|
|
return content, method
|
|
|
|
return None, None
|