working with progress bars

modified:   translate_content.py

translate_content.py

@@ -6,17 +6,23 @@ from openai import OpenAI, OpenAIError
 import pandoc
 import re
 import time
-import logging
+from tqdm import tqdm
+import sys
+import tiktoken
+import argparse
 
-# Set up logging
+# Parse command-line arguments
-logging.basicConfig(
+parser = argparse.ArgumentParser(description="Translate .rst files with progress and debugging.")
-    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+parser.add_argument("--debug", action="store_true", help="Enable debug mode with streaming output.")
-)
+args = parser.parse_args()
 
 # AI client setup
 client = OpenAI(base_url="https://hermes.ai.unturf.com/v1", api_key="choose-any-value")
 MODEL = "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic"
 
+# Tokenizer setup (used only once per file)
+TOKENIZER = tiktoken.get_encoding("cl100k_base")
+
 TOP_20_LANGS = {
     "en": "English",
     "zh": "Chinese (Simplified)",
@@ -92,7 +98,7 @@ def restore_special_content(content, placeholders):
 
 
 def translate_text(text, target_lang_full, retries=3, delay=5, timeout=100):
-    """Translate text with logging, a long timeout, and error handling."""
+    """Translate text with optional streaming debug mode."""
     prompt = (
         f"Translate the following reStructuredText (.rst) content into {target_lang_full}. "
         "The input is in English and formatted as valid reStructuredText, which includes metadata "
@@ -106,46 +112,64 @@ def translate_text(text, target_lang_full, retries=3, delay=5, timeout=100):
     )
     messages = [{"role": "user", "content": prompt}]
 
-    logging.info(
+    if args.debug:
-        f"Attempting translation to {target_lang_full} (length of input: {len(text)} characters)"
+        # Streaming mode with no newline between chunks
-    )
+        for attempt in range(retries):
-
+            try:
-    for attempt in range(retries):
+                print(f"\nDEBUG: Streaming attempt {attempt + 1}/{retries} for {target_lang_full}", file=sys.stderr)
-        try:
+                print(f"DEBUG: Translation output in {target_lang_full}: ", end="", file=sys.stderr)
-            response = client.chat.completions.create(
+                stream = client.chat.completions.create(
-                model=MODEL,
+                    model=MODEL,
-                messages=messages,
+                    messages=messages,
-                temperature=0.5,
+                    temperature=0.5,
-                max_tokens=60000,
+                    max_tokens=60000,
-                timeout=timeout,  # 100 seconds timeout
+                    stream=True,
-            )
-            result = response.choices[0].message.content
-            logging.info(
-                f"Successfully translated to {target_lang_full} on attempt {attempt + 1}"
-            )
-            return result
-        except OpenAIError as e:
-            logging.error(
-                f"OpenAI API error on attempt {attempt + 1}/{retries}: {str(e)}"
-            )
-            if attempt < retries - 1:
-                logging.info(f"Retrying in {delay} seconds...")
-                time.sleep(delay)
-            else:
-                raise Exception(
-                    f"Failed to translate to {target_lang_full} after {retries} attempts: {str(e)}"
                 )
-        except Exception as e:
+                result = ""
-            logging.error(
+                for chunk in stream:
-                f"Unexpected error on attempt {attempt + 1}/{retries}: {str(e)}"
+                    if chunk.choices[0].delta.content is not None:
-            )
+                        content = chunk.choices[0].delta.content
-            if attempt < retries - 1:
+                        result += content
-                logging.info(f"Retrying in {delay} seconds...")
+                        print(content, end="", file=sys.stderr)  # No newline, continuous stream
-                time.sleep(delay)
+                        sys.stderr.flush()  # Ensure immediate output
-            else:
+                print(f"\nDEBUG: Streaming completed for {target_lang_full}", file=sys.stderr)
-                raise Exception(
+                return result
-                    f"Translation to {target_lang_full} failed after {retries} attempts: {str(e)}"
+            except OpenAIError as e:
+                print(f"\nDEBUG: OpenAI API error on attempt {attempt + 1}/{retries}: {str(e)}", file=sys.stderr)
+                if attempt < retries - 1:
+                    print(f"DEBUG: Retrying in {delay} seconds...", file=sys.stderr)
+                    time.sleep(delay)
+                else:
+                    raise Exception(f"Failed to translate to {target_lang_full} after {retries} attempts: {str(e)}")
+            except Exception as e:
+                print(f"\nDEBUG: Unexpected error on attempt {attempt + 1}/{retries}: {str(e)}", file=sys.stderr)
+                if attempt < retries - 1:
+                    print(f"DEBUG: Retrying in {delay} seconds...", file=sys.stderr)
+                    time.sleep(delay)
+                else:
+                    raise Exception(f"Translation to {target_lang_full} failed after {retries} attempts: {str(e)}")
+    else:
+        # Non-debug, blocking mode
+        for attempt in range(retries):
+            try:
+                response = client.chat.completions.create(
+                    model=MODEL,
+                    messages=messages,
+                    temperature=0.5,
+                    max_tokens=60000,
+                    timeout=timeout,
                 )
+                return response.choices[0].message.content
+            except OpenAIError as e:
+                if attempt < retries - 1:
+                    time.sleep(delay)
+                else:
+                    raise Exception(f"Failed to translate to {target_lang_full} after {retries} attempts: {str(e)}")
+            except Exception as e:
+                if attempt < retries - 1:
+                    time.sleep(delay)
+                else:
+                    raise Exception(f"Translation to {target_lang_full} failed after {retries} attempts: {str(e)}")
 
 
 def load_hashes():
@@ -164,7 +188,7 @@ def get_content_hash(content):
     return hashlib.sha256(content.encode("utf-8")).hexdigest()
 
 
-def process_file(file_path):
+def process_file(file_path, lang_pbar):
     with open(file_path, "r", encoding="utf-8") as f:
         lines = f.readlines()
         header_end = next(
@@ -173,6 +197,9 @@ def process_file(file_path):
         header = lines[:header_end]
         content = "".join(lines[header_end:])
 
+    # Count tokens for the entire .rst content once
+    total_tokens = len(TOKENIZER.encode(content))
+
     base_name = os.path.splitext(os.path.basename(file_path))[0]
     slug = base_name
     source_hash = get_content_hash(content)
@@ -185,6 +212,7 @@ def process_file(file_path):
     replaced_content, placeholders = replace_special_content(content)
 
     for lang_code, lang_full in TOP_20_LANGS.items():
+        lang_pbar.set_description(f"File: {slug}.rst | Lang: {lang_full} ({lang_code}) | Total Tokens: {total_tokens}")
         lang_dir = f"{CONTENT_DIR}/{lang_code}"
         os.makedirs(lang_dir, exist_ok=True)
 
@@ -217,16 +245,13 @@ def process_file(file_path):
             with open(md_file, "r", encoding="utf-8") as f:
                 current_md_hash = get_content_hash(f.read())
             if current_rst_hash == rst_hash and current_md_hash == md_hash:
-                logging.info(
+                lang_pbar.update(1)
-                    f"Skipping {slug} in {lang_full} ({lang_code}) - up to date"
-                )
                 continue
 
         # Translate or use original content
         if lang_code == "en":
             translated_replaced_content = replaced_content
         else:
-            logging.info(f"Translating {slug} to {lang_full} ({lang_code})")
             translated_replaced_content = translate_text(replaced_content, lang_full)
 
         # Restore special content
@@ -261,16 +286,25 @@ def process_file(file_path):
 
         # Save hashes after each language
         save_hashes(hashes)
+        lang_pbar.update(1)
 
 
 def main():
-    for root, _, files in os.walk(CONTENT_DIR):
+    rst_files = [
-        for file in files:
+        os.path.join(root, file)
-            if file.endswith(".rst") and not any(
+        for root, _, files in os.walk(CONTENT_DIR)
-                f"/{lang}/" in root for lang in TOP_20_LANGS.keys()
+        for file in files
-            ):
+        if file.endswith(".rst")
-                logging.info(f"Processing file: {os.path.join(root, file)}")
+        and not any(f"/{lang}/" in root for lang in TOP_20_LANGS.keys())
-                process_file(os.path.join(root, file))
+    ]
+
+    with tqdm(total=len(rst_files), desc="Total Files", unit="file", file=sys.stderr) as file_pbar:
+        with tqdm(total=len(TOP_20_LANGS), desc="File: N/A | Lang: N/A | Total Tokens: N/A", unit="lang", file=sys.stderr) as lang_pbar:
+            for file_path in rst_files:
+                file_pbar.set_description(f"Total Files (Current: {os.path.basename(file_path)})")
+                lang_pbar.reset()
+                process_file(file_path, lang_pbar)
+                file_pbar.update(1)
 
 
 if __name__ == "__main__":