Automatic article translation in Hugo site

translate.py

from openai import OpenAI
from tqdm import tqdm
import re
import os
import glob
import concurrent.futures
from pathlib import Path
import time

client = OpenAI(
    base_url="https://api.deepseek.com", api_key=os.environ.get("DEEPSEEK_API_KEY")
)


def get_translation_prompt(article_chunk, language="French") -> str:
    return f"""Translate the following article chunk into {language}. Leave code blocks
    unchanged, except for comments. Match the style of the English written and
    keep the style natural but terse. Translate the front matter *values* that are in English to {language} as well.
    Do not modify any Hugo partials or shortcodes. Do not modify the formatting. Your only task is to translate this article chunk.
    Return the translated chunk in a code fence in this format:

    ```markdown
    <translated article contents>
    ```

    Here is the article chunk to translate:

    ```markdown
    {article_chunk}
    ```"""


extract_markdown_re = re.compile(r"```markdown\n(.*)```", re.DOTALL)


def translate_chunk(chunk, language="French", max_retries=3, retry_delay=1):
    chunk = re.sub(
        r"\n{3,}", "\n\n", chunk
    )  # more than 2 newlines makes the model think the article is done

    for attempt in range(max_retries):
        try:
            completion = client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {
                        "role": "user",
                        "content": get_translation_prompt(chunk, language),
                    },
                ],
                max_completion_tokens=50000,
                temperature=0.6,
            )
            response = completion.choices[0].message.content
            match = extract_markdown_re.search(response)
            if match:
                return match.group(1)
            else:
                print("No match found in response")
                print("Saving failed response")
                os.makedirs("failed_responses", exist_ok=True)
                with open(f"failed_responses/{hash(chunk)}_{attempt}.txt", "w") as f:
                    f.write(response)
        except Exception as e:
            if attempt < max_retries - 1:
                print(
                    f"Error during API call: {e}, retrying ({attempt+1}/{max_retries})..."
                )
                time.sleep(retry_delay)
            else:
                print(f"Failed after {max_retries} attempts: {e}")
                return None


def split_markdown_into_chunks(content, max_chars=None):
    """
    Split markdown into logical chunks preserving code blocks, math, and front matter.
    When joined back together, the result should be identical to the original content.
    """
    # Handle empty content
    if not content:
        return []

    # Extract front matter if present
    front_matter = ""
    content_to_process = content
    if content.startswith("---"):
        end_idx = content.find("---", 3)
        if end_idx != -1:
            front_matter = content[: end_idx + 3]
            content_to_process = content[end_idx + 3 :]

    # Parse the markdown to find code blocks
    lines = content_to_process.split("\n")
    chunks = []
    current_chunk = []
    in_code_block = False
    in_math_block = False

    # If we have front matter, start with it
    if front_matter:
        current_chunk.append(front_matter)

    for line in lines:
        # Check if this line starts or ends a code block
        if line.strip().startswith("```"):
            in_code_block = not in_code_block

        # Check if this line starts or ends a math block
        if line.strip() == "$$":
            in_math_block = not in_math_block

        # Don't split in the middle of code or math blocks
        if not in_code_block and not in_math_block:
            # Start a new chunk at headers, unless current chunk is empty
            if line.startswith("#") and current_chunk:
                chunks.append("\n".join(current_chunk))
                current_chunk = []

        current_chunk.append(line)

    # Add the final chunk
    if current_chunk:
        chunks.append("\n".join(current_chunk))

    return chunks


def translate_article(article, language="French", max_retries=3, retry_delay=1):
    # Only split into chunks if article is >3k characters
    if len(article) > 3000:
        # Split the article into chunks
        chunks = split_markdown_into_chunks(article)

        # For shorter articles, use a single chunk
        if len(chunks) == 1:
            return {"chunks": [article], "is_chunked": False}

        # Return info needed to reconstruct the article
        return {"chunks": chunks, "is_chunked": True}
    else:
        # For shorter articles, translate as a single chunk
        print(f"Translating article as a single chunk ({len(article)} chars)")
        return {"chunks": [article], "is_chunked": False}


def process_translation_task(task):
    chunk = task["chunk"]
    file_path = task.get("file_path")
    max_retries = task.get("max_retries", 3)
    retry_delay = task.get("retry_delay", 1)
    chunk_idx = task.get("chunk_idx")
    language = task.get("language", "French")
    lang_code = task.get("lang_code", "fr")

    return {
        "translated_chunk": translate_chunk(chunk, language, max_retries, retry_delay),
        "file_path": file_path,
        "chunk_idx": chunk_idx,
        "lang_code": lang_code,
    }


def translate_all_md_files(content_dir="../content", languages=None, workers=50):
    """
    Translate all markdown files in the content directory to the specified languages.

    Args:
        content_dir: Path to the content directory.
        languages: List of tuples in the format [(lang_code, language_name)].
                   Default is [("fr", "French")].
        workers: Number of workers for concurrent translation.
    """
    if languages is None:
        languages = [("fr", "French")]

    # Find all markdown files in the content directory
    md_files = glob.glob(os.path.join(content_dir, "**/*.md"), recursive=True)

    # Prepare files for each language
    all_files_to_translate = {}

    for lang_code, _ in languages:
        # Only translate English files (files without a language code in the name)
        english_files = [f for f in md_files if "." not in Path(f).stem]

        # Check for files that need translation (either no translation exists or source is newer)
        files_to_translate = []
        for file_path in english_files:
            original_path = Path(file_path)
            stem = original_path.stem
            translated_file_path = original_path.with_name(f"{stem}.{lang_code}.md")

            # Translate if translation doesn't exist or original file is newer
            if not translated_file_path.exists() or (original_path.stat().st_mtime > translated_file_path.stat().st_mtime):
                files_to_translate.append(file_path)

        all_files_to_translate[lang_code] = files_to_translate
        print(
            f"Found {len(files_to_translate)} markdown files to translate to {lang_code} out of {len(md_files)} total files."
        )
    # Prepare all translation tasks
    all_tasks = []
    file_chunks_map = {}  # Maps (file_path, lang_code) to list of chunks and metadata

    for lang_code, language in languages:
        for file_path in all_files_to_translate[lang_code]:
            try:
                # Read the markdown file
                with open(file_path, "r", encoding="utf-8") as file:
                    content = file.read()

                # Get chunks to translate
                result = translate_article(content, language)

                chunks = result["chunks"]
                is_chunked = result["is_chunked"]

                # Store file metadata
                file_key = (file_path, lang_code)
                file_chunks_map[file_key] = {
                    "chunks": chunks,
                    "is_chunked": is_chunked,
                    "results": [None] * len(chunks),
                    "language": language,
                    "lang_code": lang_code,
                }

                # Create tasks for each chunk
                for i, chunk in enumerate(chunks):
                    all_tasks.append(
                        {
                            "chunk": chunk,
                            "file_path": file_path,
                            "max_retries": 3,
                            "retry_delay": 1,
                            "chunk_idx": i,
                            "total_chunks": len(chunks),
                            "language": language,
                            "lang_code": lang_code,
                        }
                    )
            except Exception as e:
                print(f"Error preparing {file_path} for {lang_code}: {e}")

    # Execute all translation tasks with a single thread pool
    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {
            executor.submit(process_translation_task, task): task for task in all_tasks
        }

        # Process results as they complete
        for future in tqdm(
            concurrent.futures.as_completed(futures),
            total=len(futures),
            desc="Translating chunks",
        ):
            task = futures[future]
            try:
                result = future.result()
                file_path = result["file_path"]
                chunk_idx = result["chunk_idx"]
                translated_chunk = result["translated_chunk"]
                lang_code = result["lang_code"]
                file_key = (file_path, lang_code)

                if translated_chunk is None:
                    print(f"Translation to {lang_code} failed for chunk {chunk_idx+1} of {file_path}")
                    continue

                # Store the result
                file_chunks_map[file_key]["results"][chunk_idx] = translated_chunk

                # Check if all chunks for this file are done
                if all(r is not None for r in file_chunks_map[file_key]["results"]):
                    # All chunks for this file are translated, write the file
                    translated_content = "\n".join(
                        file_chunks_map[file_key]["results"]
                    )
                    original_path = Path(file_path)
                    translated_file_path = original_path.with_name(
                        f"{original_path.stem}.{lang_code}.md"
                    )

                    with open(translated_file_path, "w", encoding="utf-8") as trans_file:
                        trans_file.write(translated_content)

            except Exception as e:
                lang_code = task.get("lang_code", "unknown")
                print(
                    f"Error processing {lang_code} chunk {task['chunk_idx']+1} of {task['file_path']}: {e}"
                )

    # Count successful translations for each language
    success_counts = {}
    for lang_code, _ in languages:
        success_counts[lang_code] = sum(
            1
            for file_key, data in file_chunks_map.items()
            if file_key[1] == lang_code and all(r is not None for r in data["results"])
        )


def delete_translated_md_files(content_dir="../content", lang_code="fr"):
    translated_files = glob.glob(os.path.join(content_dir, f"**/*.{lang_code}.md"), recursive=True)

    deleted_count = 0
    for file_path in translated_files:
        try:
            os.remove(file_path)
            print(f"Deleted: {file_path}")
            deleted_count += 1
        except Exception as e:
            print(f"Error deleting {file_path}: {e}")

    print(
        f"Deletion completed: {deleted_count} of {len(translated_files)} {lang_code} translation files deleted."
    )


def translate_directory(content_dir="../content", languages=None, workers=200):
    if languages is None:
        languages = [("fr", "French")]

    translate_all_md_files(content_dir, languages, workers)


if __name__ == "__main__":
    # Example usage:
    # translate_directory(languages=[("fr", "French"), ("es", "Spanish"), ("de", "German")])
    languages = [
        ('fr', 'French'),
        ('zh', 'Chinese'),
    ]
    translate_directory(languages=languages)