Update soni_translate/translate_segments.py

soni_translate/translate_segments.py

@@ -7,6 +7,8 @@ from .logging_setup import logger
 import re
 import json
 import time
+import os
+import google.generativeai as genai
 
 TRANSLATION_PROCESS_OPTIONS = [
     "google_translator_batch",
@@ -15,12 +17,15 @@ TRANSLATION_PROCESS_OPTIONS = [
     "gpt-3.5-turbo-0125",
     "gpt-4-turbo-preview_batch",
     "gpt-4-turbo-preview",
+    "gemini-pro",
+    "gemini-pro_batch",
     "disable_translation",
 ]
 DOCS_TRANSLATION_PROCESS_OPTIONS = [
     "google_translator",
     "gpt-3.5-turbo-0125",
     "gpt-4-turbo-preview",
+    "gemini-pro",
     "disable_translation",
 ]
 
@@ -418,6 +423,74 @@ def gpt_batch(segments, model, target, token_batch_limit=900, source=None):
     )
 
 
+def check_gemini_api_key():
+    """Check if Gemini API key is set in environment variables."""
+    if not os.environ.get("GOOGLE_API_KEY"):
+        raise ValueError(
+            "Gemini API key not found. Please set the GOOGLE_API_KEY environment variable."
+        )
+
+def translate_with_gemini(text, target_lang, source_lang=None):
+    """Translate text using Google's Gemini API."""
+    check_gemini_api_key()
+    genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
+    model = genai.GenerativeModel('gemini-pro')
+    
+    prompt = f"""Translate the following text to {target_lang}. 
+    Keep the same tone and style. Preserve any special characters or formatting.
+    
+    Text to translate: {text}
+    """
+    if source_lang:
+        prompt = f"Translate from {source_lang} to {target_lang}: {text}"
+        
+    response = model.generate_content(prompt)
+    return response.text.strip()
+
+def gemini_sequential(segments, target, source=None):
+    """Translate segments sequentially using Gemini."""
+    segments_ = copy.deepcopy(segments)
+    
+    for line in tqdm(range(len(segments_))):
+        text = segments_[line]["text"]
+        translated_line = translate_with_gemini(text.strip(), target, source)
+        segments_[line]["text"] = translated_line
+        
+    return segments_
+
+def gemini_batch(segments, target, token_batch_limit=1000, source=None):
+    """Translate segments in batches using Gemini."""
+    segments_ = copy.deepcopy(segments)
+    batch_texts = []
+    current_batch = []
+    current_length = 0
+    
+    # Group texts into batches
+    for segment in segments_:
+        text_length = len(segment["text"])
+        if current_length + text_length > token_batch_limit:
+            batch_texts.append(current_batch)
+            current_batch = []
+            current_length = 0
+        current_batch.append(segment["text"])
+        current_length += text_length
+    
+    if current_batch:
+        batch_texts.append(current_batch)
+    
+    # Translate each batch
+    for i, batch in enumerate(tqdm(batch_texts)):
+        batch_text = "\n---\n".join(batch)
+        translated_batch = translate_with_gemini(batch_text, target, source)
+        translated_segments = translated_batch.split("\n---\n")
+        
+        # Update segments with translations
+        start_idx = sum(len(b) for b in batch_texts[:i])
+        for j, translation in enumerate(translated_segments):
+            segments_[start_idx + j]["text"] = translation.strip()
+    
+    return segments_
+
 def translate_text(
     segments,
     target,
@@ -443,7 +516,7 @@ def translate_text(
             )
         case model if model in ["gpt-3.5-turbo-0125", "gpt-4-turbo-preview"]:
             return gpt_sequential(segments, model, target, source)
-        case model if model in ["gpt-3.5-turbo-0125_batch", "gpt-4-turbo-preview_batch",]:
+        case model if model in ["gpt-3.5-turbo-0125_batch", "gpt-4-turbo-preview_batch"]:
             return gpt_batch(
                 segments,
                 translation_process.replace("_batch", ""),
@@ -451,6 +524,10 @@ def translate_text(
                 token_batch_limit,
                 source
             )
+        case "gemini-pro":
+            return gemini_sequential(segments, target, source)
+        case "gemini-pro_batch":
+            return gemini_batch(segments, target, token_batch_limit, source)
         case "disable_translation":
             return segments
         case _: