Update ycecream loguru, add tz

app.py

@@ -1,53 +1,85 @@
-import gradio as gr
+import os
 import re
+import time
 from difflib import Differ
-from src.translation_agent.utils import *
+
+import gradio as gr
+from loguru import logger
+from ycecream import y
+
+# from src.translation_agent.utils import *
+from src.translation_agent.utils import MAX_TOKENS_PER_CHUNK
+from src.translation_agent.utils import RecursiveCharacterTextSplitter
+from src.translation_agent.utils import calculate_chunk_size
+from src.translation_agent.utils import multichunk_improve_translation
+from src.translation_agent.utils import multichunk_initial_translation
+from src.translation_agent.utils import multichunk_reflect_on_translation
+from src.translation_agent.utils import num_tokens_in_string
+from src.translation_agent.utils import one_chunk_improve_translation
+from src.translation_agent.utils import one_chunk_initial_translation
+from src.translation_agent.utils import one_chunk_reflect_on_translation
+
+y.configure(sln=1, show_time=1)
+
+os.environ["TZ"] = "Asia/Shanghai"
+try:
+    time.tzset()  # type: ignore # pylint: disable=no-member
+except Exception:
+    # Windows
+    y("Windows, cant run time.tzset()")
 
 LANGUAGES = {
-    'English': 'English',
-    'Español': 'Spanish',
-    'Français': 'French',
-    'Deutsch': 'German',
-    'Italiano': 'Italian',
-    'Português': 'Portuguese',
-    'Русский': 'Russian',
-    '中文': 'Chinese',
-    '日本語': 'Japanese',
-    '한국어': 'Korean',
-    'العربية': 'Arabic',
-    'हिन्दी': 'Hindi',
+    "English": "English",
+    "Español": "Spanish",
+    "Français": "French",
+    "Deutsch": "German",
+    "Italiano": "Italian",
+    "Português": "Portuguese",
+    "Русский": "Russian",
+    "中文": "Chinese",
+    "日本語": "Japanese",
+    "한국어": "Korean",
+    "العربية": "Arabic",
+    "हिन्दी": "Hindi",
 }
 
 
 def diff_texts(text1, text2, lang):
     d = Differ()
-    ic(lang)
-    if lang == '中文':
+    # y(lang)
+    logger.info(lang)
+    if lang == "中文":
         return [
             (token[2:], token[0] if token[0] != " " else None)
             for token in d.compare(text1, text2)
             if token[0] in ["+", " "]
         ]
     else:
-        words1 = re.findall(r'\S+|\s+', text1)
-        words2 = re.findall(r'\S+|\s+', text2)
+        words1 = re.findall(r"\S+|\s+", text1)
+        words2 = re.findall(r"\S+|\s+", text2)
 
         return [
             (token[2:], token[0] if token[0] != " " else None)
             for token in d.compare(words1, words2)
             if token[0] in ["+", " "]
         ]
-    
 
-def translate_text(source_lang, target_lang, source_text, country, max_tokens=MAX_TOKENS_PER_CHUNK):
+
+def translate_text(
+    source_lang,
+    target_lang,
+    source_text,
+    country,
+    max_tokens=MAX_TOKENS_PER_CHUNK,
+):
     num_tokens_in_text = num_tokens_in_string(source_text)
 
-    ic(num_tokens_in_text)
+    y(num_tokens_in_text)
 
     if num_tokens_in_text < max_tokens:
-        ic("Translating text as single chunk")
+        y("Translating text as single chunk")
 
-        #Note: use yield from B() if put yield in function B()
+        # Note: use yield from B() if put yield in function B()
         translation_1 = one_chunk_initial_translation(
             source_lang, target_lang, source_text
         )
@@ -61,30 +93,32 @@ def translate_text(source_lang, target_lang, source_text, country, max_tokens=MA
         translation_2 = one_chunk_improve_translation(
             source_lang, target_lang, source_text, translation_1, reflection
         )
-        translation_diff = diff_texts(translation_1, translation_2, target_lang)
+        translation_diff = diff_texts(
+            translation_1, translation_2, target_lang
+        )
         yield translation_1, reflection, translation_diff
 
     else:
-        ic("Translating text as multiple chunks")
+        y("Translating text as multiple chunks")
 
         token_size = calculate_chunk_size(
             token_count=num_tokens_in_text, token_limit=max_tokens
         )
 
-        ic(token_size)
+        y(token_size)
 
         text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
-            model_name = "gpt-4",
+            model_name="gpt-4",
             chunk_size=token_size,
             chunk_overlap=0,
         )
 
         source_text_chunks = text_splitter.split_text(source_text)
-        
+
         translation_1_chunks = multichunk_initial_translation(
             source_lang, target_lang, source_text_chunks
         )
-        ic(translation_1_chunks)
+        y(translation_1_chunks)
         translation_1 = "".join(translation_1_chunks)
         yield translation_1, None, None
 
@@ -95,7 +129,7 @@ def translate_text(source_lang, target_lang, source_text, country, max_tokens=MA
             translation_1_chunks,
             country,
         )
-        ic(reflection_chunks)
+        y(reflection_chunks)
         reflection = "".join(reflection_chunks)
         yield translation_1, reflection, None
 
@@ -106,24 +140,40 @@ def translate_text(source_lang, target_lang, source_text, country, max_tokens=MA
             translation_1_chunks,
             reflection_chunks,
         )
-        ic(translation_2_chunks)
+        y(translation_2_chunks)
         translation_2 = "".join(translation_2_chunks)
-        translation_diff = diff_texts(translation_1, translation_2, target_lang)
-        
+        translation_diff = diff_texts(
+            translation_1, translation_2, target_lang
+        )
+
         yield translation_1, reflection, translation_diff
-        
 
 
 def update_ui(translation_1, reflection, translation_diff):
-    return gr.update(value=translation_1), gr.update(value=reflection), gr.update(value=translation_diff)
+    return (
+        gr.update(value=translation_1),
+        gr.update(value=reflection),
+        gr.update(value=translation_diff),
+    )
+
 
 with gr.Blocks() as demo:
     gr.Markdown("# Andrew Ng's Translation Agent ")
     with gr.Row():
-        source_lang = gr.Dropdown(choices=list(LANGUAGES.keys()), value='English', label="Source Language")
-        target_lang = gr.Dropdown(choices=list(LANGUAGES.keys()), value='中文', label="Target Language")
+        source_lang = gr.Dropdown(
+            choices=list(LANGUAGES.keys()),
+            value="English",
+            label="Source Language",
+        )
+        target_lang = gr.Dropdown(
+            choices=list(LANGUAGES.keys()),
+            value="中文",
+            label="Target Language",
+        )
         country = gr.Textbox(label="Country (for target language)")
-    source_text = gr.Textbox(label="Source Text", lines=5, show_copy_button=True)
+    source_text = gr.Textbox(
+        label="Source Text", lines=5, show_copy_button=True
+    )
 
     btn = gr.Button("Translate")
 
@@ -131,13 +181,25 @@ with gr.Blocks() as demo:
         translation_1 = gr.Textbox(label="Initial Translation", lines=3)
         reflection = gr.Textbox(label="Reflection", lines=3)
 
-    translation_diff = gr.HighlightedText  (label="Final Translation",
-                                            combine_adjacent=True,
-                                            show_legend=True,
-                                            color_map={"+": "red"})
-    #translation = gr.Textbox(label="Final Translation", lines=5, show_copy_button=True)
-
-    btn.click(translate_text, inputs=[source_lang, target_lang, source_text, country], outputs=[translation_1, reflection, translation_diff], queue=True)
-    btn.click(update_ui, inputs=[translation_1, reflection, translation_diff], outputs=[translation_1, reflection, translation_diff], queue=True)
+    translation_diff = gr.HighlightedText(
+        label="Final Translation",
+        combine_adjacent=True,
+        show_legend=True,
+        color_map={"+": "red"},
+    )
+    # translation = gr.Textbox(label="Final Translation", lines=5, show_copy_button=True)
+
+    btn.click(
+        translate_text,
+        inputs=[source_lang, target_lang, source_text, country],
+        outputs=[translation_1, reflection, translation_diff],
+        queue=True,
+    )
+    btn.click(
+        update_ui,
+        inputs=[translation_1, reflection, translation_diff],
+        outputs=[translation_1, reflection, translation_diff],
+        queue=True,
+    )
 
 demo.launch()

requirements.txt

@@ -1,3 +1,5 @@
+ycecream
+loguru
 aiofiles==23.2.1 ; python_version >= "3.9" and python_version < "4.0"
 altair==5.3.0 ; python_version >= "3.9" and python_version < "4.0"
 annotated-types==0.7.0 ; python_version >= "3.9" and python_version < "4.0"