init deliver

.gitignore

@@ -0,0 +1,5 @@
+.env
+venv
+.venv
+**.pyc
+downloads/

app.py

@@ -0,0 +1,454 @@
+import gradio as gr
+import json
+import os
+import sys
+import tempfile
+import whisperx
+import ffmpeg
+import tiktoken
+from io import BytesIO
+from glob import glob
+from dotenv import load_dotenv
+from download import download_video_audio, delete_download
+from groq import Groq
+from openai import OpenAI
+from langdetect import detect
+from langdetect.lang_detect_exception import LangDetectException
+from translation_agent.utils import *
+
+
+os.environ["FFMPEG_PATH"] = "D:\\ffmpeg\\bin\\ffmpeg.exe"
+os.environ["no_proxy"] = "localhost, 127.0.0.1, ::1"
+
+load_dotenv()
+
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY", None)
+client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL"))
+model = os.getenv("OPENAI_MODEL") or "gpt-4o"
+
+MAX_FILE_SIZE = 100 * 1024 * 1024  # 100 MB
+FILE_TOO_LARGE_MESSAGE = "The audio file is too large. If you used a YouTube link, please try a shorter video clip. If you uploaded an audio file, try trimming or compressing the audio to under 100 MB."
+
+audio_file_path = None
+progress=gr.Progress()
+
+
+def detect_language(text):
+    try:
+        language = detect(text)
+        return language
+    except LangDetectException as e:
+        print(f"Error detecting language: {e}")
+        return None
+    
+class GenerationStatistics:
+    def __init__(self, input_time=0, output_time=0, input_tokens=0, output_tokens=0, total_time=0, model_name=model):
+        self.input_time = input_time
+        self.output_time = output_time
+        self.input_tokens = input_tokens
+        self.output_tokens = output_tokens
+        self.total_time = total_time
+        self.model_name = model_name
+
+    def get_input_speed(self):
+        return self.input_tokens / self.input_time if self.input_time != 0 else 0
+    
+    def get_output_speed(self):
+        return self.output_tokens / self.output_time if self.output_time != 0 else 0
+    
+    def add(self, other):
+        if not isinstance(other, GenerationStatistics):
+            raise TypeError("Can only add GenerationStatistics objects")
+        
+        self.input_time += other.input_time
+        self.output_time += other.output_time
+        self.input_tokens += other.input_tokens
+        self.output_tokens += other.output_tokens
+        self.total_time += other.total_time
+
+    def __str__(self):
+        return (f"\n## {self.get_output_speed():.2f} T/s ⚡\nRound trip time: {self.total_time:.2f}s  Model: {self.model_name}\n\n"
+                f"| Metric          | Input          | Output          | Total          |\n"
+                f"|-----------------|----------------|-----------------|----------------|\n"
+                f"| Speed (T/s)     | {self.get_input_speed():.2f}            | {self.get_output_speed():.2f}            | {(self.input_tokens + self.output_tokens) / self.total_time if self.total_time != 0 else 0:.2f}            |\n"
+                f"| Tokens          | {self.input_tokens}            | {self.output_tokens}            | {self.input_tokens + self.output_tokens}            |\n"
+                f"| Inference Time (s) | {self.input_time:.2f}            | {self.output_time:.2f}            | {self.total_time:.2f}            |")
+
+class NoteSection:
+    def __init__(self, structure, transcript):
+        self.structure = structure
+        self.contents = {title: "" for title in self.flatten_structure(structure)}
+
+    def flatten_structure(self, structure):
+        sections = []
+        for title, content in structure.items():
+            sections.append(title)
+            if isinstance(content, dict):
+                sections.extend(self.flatten_structure(content))
+        return sections
+
+    def update_content(self, title, new_content):
+        try:
+            self.contents[title] += new_content
+        except TypeError as e:
+            pass
+
+    def return_existing_contents(self, level=1) -> str:
+        existing_content = ""
+        for title, content in self.structure.items():
+            if self.contents[title].strip():
+                existing_content += f"{'#' * level} {title}\n{self.contents[title]}.\n\n"
+            if isinstance(content, dict):
+                existing_content += self.get_markdown_content(content, level + 1)
+        return existing_content
+
+    def get_markdown_content(self, structure=None, level=1):
+        if structure is None:
+            structure = self.structure
+        
+        markdown_content = ""
+        for title, content in structure.items():
+            if self.contents[title].strip():
+                markdown_content += f"{'#' * level} {title}\n{self.contents[title]}.\n\n"
+            if isinstance(content, dict):
+                markdown_content += self.get_markdown_content(content, level + 1)
+        return markdown_content
+
+def transcribe_audio(audio_file):
+    model = whisperx.load_model("small", device="cuda", compute_type="float16")
+    result = model.transcribe(audio_file)
+
+    transcription = ''
+    segments = result['segments']
+    for segment in segments:
+        transcription += segment['text']
+        transcription += '\n'
+
+    return transcription
+
+def generate_notes_structure(transcript: str, model: str = model, lang: str="en"):    
+    
+    shot_example = """
+"Introduction": "Introduction to the AMA session, including the topic of Groq scaling architecture and the panelists",
+"Panelist Introductions": "Brief introductions from Igor, Andrew, and Omar, covering their backgrounds and roles at Groq",
+"Groq Scaling Architecture Overview": "High-level overview of Groq's scaling architecture, covering hardware, software, and cloud components",
+"Hardware Perspective": "Igor's overview of Groq's hardware approach, using an analogy of city traffic management to explain the traditional compute approach and Groq's innovative approach",
+"Traditional Compute": "Description of traditional compute approach, including asynchronous nature, queues, and poor utilization of infrastructure",
+"Groq's Approach": "Description of Groq's approach, including pre-orchestrated movement of data, low latency, high energy efficiency, and high utilization of resources",
+"Hardware Implementation": "Igor's explanation of the hardware implementation, including a comparison of GPU and LPU architectures"
+}"""
+
+    messages=[
+        {
+            "role": "system",
+            "content": "Write in JSON format:\n\n{\"Title of section goes here\":\"Description of section goes here\",\"Title of section goes here\":\"Description of section goes here\",\"Title of section goes here\":\"Description of section goes here\"}"
+        },
+        {
+            "role": "user",
+            #"content": f"Please respond in language [{lang}]. ### Transcript {transcript}\n\n### Example\n\n{shot_example}### Instructions\n\nCreate a structure for comprehensive notes on the above Transcript information. The output format is shown in examples. The section titles and content descriptions must be comprehensive and adapt to the content of transcript. Quality over quantity. "
+            "content": f"Please respond in language [{lang}]. ### Transcript {transcript}\n\n### Instructions\n\nCreate a structure for comprehensive notes on the above Transcript information. The output is json format as shown in system content. Section titles and content descriptions must be comprehensive and adapt to the content of transcript. Quality over quantity. "
+        }
+    ]
+
+    completion = client.chat.completions.create(
+        model=model,
+        messages=messages,
+        temperature=0.3,
+        #max_tokens=128000,
+        top_p=1,
+        stream=False,
+        stop=None,
+    )
+
+    usage = completion.usage
+    #statistics_to_return = GenerationStatistics(input_tokens=usage.prompt_tokens, output_tokens=usage.completion_tokens, model_name=model)
+    return completion.choices[0].message.content
+
+def generate_section(transcript: str, existing_notes: str, section: str, model: str = model, lang: str="en"):
+    
+    stream = client.chat.completions.create(
+        model=model,
+        messages=[
+            {
+                "role": "system",
+                "content": "Respond in language [{lang}]. You are an expert writer. Generate a comprehensive note for the section provided based factually on the transcript provided. Do *not* repeat any content from previous sections."
+            },
+            {
+                "role": "user",
+                "content": f"Respond in language [{lang}]. ### Transcript\n\n{transcript}\n\n### Existing Notes\n\n{existing_notes}\n\n### Instructions\n\nGenerate comprehensive notes for this section only based on the transcript: \n\n{section}."
+            }
+        ],
+        temperature=0.3,
+        max_tokens=16000,
+        top_p=1,
+        stream=True,
+        stop=None,
+    )
+
+    for chunk in stream:
+        tokens = chunk.choices[0].delta.content
+        if tokens:
+            yield tokens
+
+def process_audio(audio_file, youtube_link):
+    global audio_file_path
+    
+    i = 1; j = 3    #for progress
+    if youtube_link:
+        j += 1; progress((i,j), desc="download video..."); i += 1
+        audio_file_path = download_video_audio(youtube_link)
+        if audio_file_path is None:
+            return "Failed to download audio from YouTube link. Please try again."
+    elif audio_file:
+        audio_file_path = audio_file.name
+    else:
+        return "Please provide either an audio file or a YouTube link."
+
+    if os.path.getsize(audio_file_path) > MAX_FILE_SIZE:
+        return FILE_TOO_LARGE_MESSAGE
+
+    progress((i,j), desc="Start transcribe audio..."); i += 1
+    transcription_text = transcribe_audio(audio_file_path)
+    print("transcription_text=",transcription_text)
+    encoding = tiktoken.get_encoding("cl100k_base")
+    token_count = len(encoding.encode(transcription_text))
+    print("token_count=",token_count)
+    #transcription_output.label = f"tokens: {token_count}"
+    lang = detect_language(transcription_text[:100])
+    print("detect lang=",lang)
+
+    progress((i,j), desc="Generating notes structure..."); i += 1
+    notes_structure = generate_notes_structure(transcription_text, lang=lang)
+    print("notes_structure=", notes_structure)
+
+    progress((i,j), desc="Generating notes section..."); i += 1
+    try:
+        notes_structure_json = json.loads(notes_structure)
+        print("notes_structure_json=",notes_structure_json)
+        notes = NoteSection(structure=notes_structure_json, transcript=transcription_text)
+
+        total_generation_statistics = GenerationStatistics(model_name=model)
+
+        for title, content in notes_structure_json.items():
+            if isinstance(content, str):
+                content_stream = generate_section(transcript=transcription_text, existing_notes=notes.return_existing_contents(), section=(title + ": " + content), lang=lang)
+                for chunk in content_stream:
+                    if isinstance(chunk, GenerationStatistics):
+                        total_generation_statistics.add(chunk)
+                    elif chunk is not None:
+                        notes.update_content(title, chunk)
+
+        return transcription_text, notes.get_markdown_content(), notes.get_markdown_content()
+    except json.JSONDecodeError:
+        return "Failed to decode the notes structure. Please try again."
+
+def generate_notes(transcription_text):
+    lang = detect_language(transcription_text[:100])
+    print("detect lang=",lang)
+    notes_structure = generate_notes_structure(transcription_text, lang=lang)
+    print("notes_structure=", notes_structure)
+
+    try:
+        notes_structure_json = json.loads(notes_structure)
+        print("notes_structure_json=",notes_structure_json)
+        notes = NoteSection(structure=notes_structure_json, transcript=transcription_text)
+
+        total_generation_statistics = GenerationStatistics(model_name=model)
+
+        for title, content in notes_structure_json.items():
+            if isinstance(content, str):
+                content_stream = generate_section(transcript=transcription_text, existing_notes=notes.return_existing_contents(), section=(title + ": " + content), lang=lang)
+                for chunk in content_stream:
+                    if isinstance(chunk, GenerationStatistics):
+                        total_generation_statistics.add(chunk)
+                    elif chunk is not None:
+                        notes.update_content(title, chunk)
+
+        return notes.get_markdown_content(), notes.get_markdown_content()
+    except json.JSONDecodeError:
+        return "Failed to decode the notes structure. Please try again."
+'''
+iface = gr.Interface(
+    fn=process_audio,
+    inputs=[
+        gr.File(label="Upload Audio File"),
+        gr.Textbox(label="YouTube Link"),
+        gr.Textbox(label="Groq API Key", type="password")
+    ],
+    outputs=gr.Textbox(label="Generated Notes"),
+    title="GroqNotes: Create structured notes from audio",
+    description="Generate notes from audio using Groq, Whisper, and Llama3"
+)
+'''
+CSS = """
+    h1 {
+        text-align: center;
+        display: block;
+        height: 10vh;
+        align-content: center;
+    }
+    footer {
+        visibility: hidden;
+    }
+    .texts {
+        min-height: 100px;
+    }
+"""
+
+def clear():
+    return None, None, None, None, None, None
+
+
+def translate_text(source_text, source_lang, target_lang,  country=None, max_tokens=MAX_TOKENS_PER_CHUNK):
+    ic(f"start to translate transcription from {source_lang} to {target_lang}")
+  
+    num_tokens_in_text = num_tokens_in_string(source_text)
+    ic(num_tokens_in_text)
+
+    if num_tokens_in_text < max_tokens:
+        ic("Translating text as single chunk")
+
+        progress((1,3), desc="First translation...")
+        #Note: use yield from B() if put yield in function B()
+        translation_1 = one_chunk_initial_translation(
+            source_lang, target_lang, source_text
+        )
+        yield translation_1, None
+
+        progress((2,3), desc="Reflecton...")
+        reflection = one_chunk_reflect_on_translation(
+            source_lang, target_lang, source_text, translation_1, country
+        )
+        yield reflection, None
+
+        progress((3,3), desc="Final translation...")
+        translation_2 = one_chunk_improve_translation(
+            source_lang, target_lang, source_text, translation_1, reflection
+        )
+        
+        yield translation_2, None
+
+    else:
+        ic("Translating text as multiple chunks")
+
+        token_size = calculate_chunk_size(
+            token_count=num_tokens_in_text, token_limit=max_tokens
+        )
+
+        ic(token_size)
+
+        text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+            model_name = "gpt-4",
+            chunk_size=token_size,
+            chunk_overlap=0,
+        )
+
+        source_text_chunks = text_splitter.split_text(source_text)
+        progress((1,3), desc="First translation...")
+        translation_1_chunks = multichunk_initial_translation(
+            source_lang, target_lang, source_text_chunks
+        )
+        ic(translation_1_chunks)
+        translation_1 = "".join(translation_1_chunks)
+        #yield translation_1, None, None
+
+        progress((2,3), desc="Reflecton...")
+        reflection_chunks = multichunk_reflect_on_translation(
+            source_lang,
+            target_lang,
+            source_text_chunks,
+            translation_1_chunks,
+            country,
+        )
+        ic(reflection_chunks)
+        reflection = "".join(reflection_chunks)
+        #yield translation_1, reflection, None
+        progress((3,3), desc="Final translation...")
+        translation_2_chunks = multichunk_improve_translation(
+            source_lang,
+            target_lang,
+            source_text_chunks,
+            translation_1_chunks,
+            reflection_chunks,
+        )
+        ic(translation_2_chunks)
+        translation_2 = "".join(translation_2_chunks)  
+        yield translation_2, None
+    
+def export_txt(strings):
+    if (strings is not None):
+        os.makedirs("outputs", exist_ok=True)
+        base_count = len(glob(os.path.join("outputs", "*.txt")))
+        file_path = os.path.join("outputs", f"{base_count:06d}.txt")
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(strings)
+        return gr.update(value=file_path, label = "Ready to download markdown Summary")
+
+def update_ui(transcription_output):
+    tokens = ""
+    encoding = tiktoken.get_encoding("cl100k_base")
+    token_count = len(encoding.encode(transcription_output))
+    print("token_count=",token_count)
+    if (token_count > 0):
+        tokens = f"tokens: {token_count}"
+    return gr.update(label=tokens)
+
+with gr.Blocks(theme="soft", css=CSS) as demo:
+    gr.Markdown("# Whisper and Translation Agent")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            endpoint = gr.Dropdown(
+                label="Endpoint",
+                choices=["Groq","OpenAI","DeepSeek","Baichuan","Ollama","Huggingface"],
+                value="Groq",
+            )
+            model = gr.Textbox(label="Model", value="llama3-70b-8192", )
+            api_key = gr.Textbox(label="API_KEY", type="password", )
+
+        with gr.Column(scale=5):              
+            with gr.Row():
+                file_input = gr.File(file_types=["audio", "video"])
+                text_input = gr.Textbox(placeholder="Enter youtube link")
+            
+            with gr.Row():
+                clear_btn = gr.Button("CLEAR")
+                extract_btn = gr.Button("Extract")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            source_lang = gr.Dropdown(label="Source Lang(Auto-Detect)", choices=["English", "Chinese", "Spanish"], value="English")
+            target_lang = gr.Dropdown(label="Target Lang", choices=["English", "Chinese", "Spanish"], value="Chinese")
+            switch_Btn = gr.Button(value="🔄️")
+            translate_btn = gr.Button("Translate")
+            download_btn = gr.DownloadButton(label="Download")
+    
+
+        with gr.Column(scale=5):
+            with gr.Row():
+                with gr.Tab("Transcription"):
+                    transcription_output = gr.Textbox(label='', lines=5, show_copy_button=True)
+                                
+                with gr.Tab("Summary"):
+                    summary_output = gr.Textbox(label='', lines=5, show_copy_button=True, elem_classes="texts")
+            
+                with gr.Tab("Markdown"):
+                    markdown_output = gr.Markdown(label='Markdown Summary', elem_classes="texts", height=500)  
+    
+            with gr.Row():
+                with gr.Tab("Translated Transcription"):
+                    translated_transcription_output = gr.Textbox(label='', lines=5, show_copy_button=True)
+                with gr.Tab("Translated Summary"):
+                    translated_summary_output = gr.Textbox(label='', elem_classes="texts", lines=5, show_copy_button=True)
+                with gr.Tab("Markdown Summary"):
+                    translated_summary_markdown = gr.Markdown(label='', elem_classes="texts", height=500)                    
+    
+    clear_btn.click(clear, outputs=[file_input, text_input, transcription_output, summary_output, translated_transcription_output, translated_summary_output])
+    extract_btn.click(process_audio, inputs=[file_input, text_input], outputs=[transcription_output, summary_output, markdown_output])
+    translate_btn.click(translate_text, inputs=[transcription_output, source_lang, target_lang], outputs=[translated_transcription_output]
+                        ).then(
+                        generate_notes, inputs=[translated_transcription_output], outputs=[translated_summary_output,translated_summary_markdown])
+    translated_summary_markdown.change(fn=export_txt, inputs=translated_summary_markdown, outputs=[download_btn])
+
+if __name__ == "__main__":
+    demo.launch()

download.py

@@ -0,0 +1,91 @@
+from __future__ import unicode_literals
+import yt_dlp as youtube_dl
+import os
+import time
+import os
+import shutil
+
+max_retries = 3
+delay = 2
+
+
+class MyLogger(object):
+    def __init__(self, external_logger=lambda x: None):
+        self.external_logger = external_logger
+
+    def debug(self, msg):
+        print("[debug]: ", msg)
+        self.external_logger(msg)
+
+    def warning(self, msg):
+        print("[warning]: ", msg)
+
+    def error(self, msg):
+        print("[error]: ", msg)
+
+
+def my_hook(d):
+    print("hook", d["status"])
+    if d["status"] == "finished":
+        print("Done downloading, now converting ...")
+
+
+def get_ydl_opts(external_logger=lambda x: None):
+    return {
+    "format": "bestaudio/best",
+    "postprocessors": [
+        {
+            "key": "FFmpegExtractAudio",
+            "preferredcodec": "mp3",
+                "preferredquality": "192",  # set the preferred bitrate to 192kbps
+            }
+        ],
+        "logger": MyLogger(external_logger),
+    "outtmpl": "./downloads/audio/%(title)s.%(ext)s",  # Set the output filename directly
+    "progress_hooks": [my_hook],
+}
+
+
+def download_video_audio(url, external_logger=lambda x: None):
+    retries = 0
+    while retries < max_retries:
+        try:
+            ydl_opts = get_ydl_opts(external_logger)
+            with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+                print("Going to download ", url)
+                info = ydl.extract_info(url, download=False)
+                filename = ydl.prepare_filename(info)
+                res = ydl.download([url])
+                print("youtube-dl result :", res)
+                mp3_filename = os.path.splitext(filename)[0] + '.mp3'
+                print('mp3 file name - ', mp3_filename)
+                return mp3_filename
+        except Exception as e:
+            retries += 1
+            print(
+                f"An error occurred during downloading (Attempt {retries}/{max_retries}):",
+                str(e),
+            )
+            if retries >= max_retries:
+                raise e
+            time.sleep(delay)
+
+
+
+def delete_download(path):
+    try:
+        if os.path.isfile(path):
+            os.remove(path)
+            print(f"File {path} has been deleted.")
+        elif os.path.isdir(path):
+            shutil.rmtree(path)
+            print(f"Directory {path} and its contents have been deleted.")
+        else:
+            print(f"The path {path} is neither a file nor a directory.")
+    except PermissionError:
+        print(f"Permission denied: Unable to delete {path}.")
+    except FileNotFoundError:
+        print(f"File or directory not found: {path}")
+    except Exception as e:
+        print(f"An error occurred while trying to delete {path}: {str(e)}")
+        

requirements.txt

@@ -0,0 +1,11 @@
+ffmpeg_python==0.2.0
+gradio==4.38.1
+groq==0.9.0
+icecream==2.1.3
+langchain_text_splitters==0.2.2
+langdetect==1.0.9
+openai==1.35.14
+python-dotenv==1.0.1
+tiktoken==0.5.2
+whisperx==3.1.1
+yt_dlp==2024.7.16

translation_agent/__init__.py

@@ -0,0 +1 @@
+from .utils import translate

translation_agent/utils.py

@@ -0,0 +1,689 @@
+import os
+from typing import List
+from typing import Union
+
+import openai
+import tiktoken
+from dotenv import load_dotenv
+from icecream import ic
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+
+load_dotenv()  # read local .env file
+model = os.getenv("OPENAI_MODEL") or "gpt-4-turbo"
+client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"), base_url=os.getenv("OPENAI_BASE_URL"))
+
+
+MAX_TOKENS_PER_CHUNK = (
+    1000  # if text is more than this many tokens, we'll break it up into
+)
+# discrete chunks to translate one chunk at a time
+
+
+def get_completion(
+    prompt: str,
+    system_message: str = "You are a helpful assistant.",
+    model: str = model,
+    temperature: float = 0.3,
+    json_mode: bool = False,
+) -> Union[str, dict]:
+    """
+        Generate a completion using the OpenAI API.
+
+    Args:
+        prompt (str): The user's prompt or query.
+        system_message (str, optional): The system message to set the context for the assistant.
+            Defaults to "You are a helpful assistant.".
+        model (str, optional): The name of the OpenAI model to use for generating the completion.
+            Defaults to "gpt-4-turbo".
+        temperature (float, optional): The sampling temperature for controlling the randomness of the generated text.
+            Defaults to 0.3.
+        json_mode (bool, optional): Whether to return the response in JSON format.
+            Defaults to False.
+
+    Returns:
+        Union[str, dict]: The generated completion.
+            If json_mode is True, returns the complete API response as a dictionary.
+            If json_mode is False, returns the generated text as a string.
+    """
+
+    if json_mode:
+        response = client.chat.completions.create(
+            model=model,
+            temperature=temperature,
+            top_p=1,
+            response_format={"type": "json_object"},
+            messages=[
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": prompt},
+            ],
+        )
+        return response.choices[0].message.content
+    else:
+        response = client.chat.completions.create(
+            model=model,
+            temperature=temperature,
+            top_p=1,
+            messages=[
+                {"role": "system", "content": system_message},
+                {"role": "user", "content": prompt},
+            ],
+        )
+        return response.choices[0].message.content
+
+
+def one_chunk_initial_translation(
+    source_lang: str, target_lang: str, source_text: str
+) -> str:
+    """
+    Translate the entire text as one chunk using an LLM.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for translation.
+        source_text (str): The text to be translated.
+
+    Returns:
+        str: The translated text.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation from {source_lang} to {target_lang}."
+
+    translation_prompt = f"""This is an {source_lang} to {target_lang} translation, please provide the {target_lang} translation for this text. \
+Do not provide any explanations or text apart from the translation.
+{source_lang}: {source_text}
+
+{target_lang}:"""
+
+    prompt = translation_prompt.format(source_text=source_text)
+
+    translation = get_completion(prompt, system_message=system_message)
+
+    return translation
+
+
+def one_chunk_reflect_on_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text: str,
+    translation_1: str,
+    country: str = "",
+) -> str:
+    """
+    Use an LLM to reflect on the translation, treating the entire text as one chunk.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language of the translation.
+        source_text (str): The original text in the source language.
+        translation_1 (str): The initial translation of the source text.
+        country (str): Country specified for target language.
+
+    Returns:
+        str: The LLM's reflection on the translation, providing constructive criticism and suggestions for improvement.
+    """
+
+    system_message = f"You are an expert linguist specializing in translation from {source_lang} to {target_lang}. \
+You will be provided with a source text and its translation and your goal is to improve the translation."
+
+    if country != "":
+        reflection_prompt = f"""Your task is to carefully read a source text and a translation from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions to improve the translation. \
+The final style and tone of the translation should match the style of {target_lang} colloquially spoken in {country}.
+
+The source text and initial translation, delimited by XML tags <SOURCE_TEXT></SOURCE_TEXT> and <TRANSLATION></TRANSLATION>, are as follows:
+
+<SOURCE_TEXT>
+{source_text}
+</SOURCE_TEXT>
+
+<TRANSLATION>
+{translation_1}
+</TRANSLATION>
+
+When writing suggestions, pay attention to whether there are ways to improve the translation's \n\
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),\n\
+(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules, and ensuring there are no unnecessary repetitions),\n\
+(iii) style (by ensuring the translations reflect the style of the source text and takes into account any cultural context),\n\
+(iv) terminology (by ensuring terminology use is consistent and reflects the source text domain; and by only ensuring you use equivalent idioms {target_lang}).\n\
+
+Write a list of specific, helpful and constructive suggestions for improving the translation.
+Each suggestion should address one specific part of the translation.
+Output only the suggestions and nothing else."""
+
+    else:
+        reflection_prompt = f"""Your task is to carefully read a source text and a translation from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions to improve the translation. \
+
+The source text and initial translation, delimited by XML tags <SOURCE_TEXT></SOURCE_TEXT> and <TRANSLATION></TRANSLATION>, are as follows:
+
+<SOURCE_TEXT>
+{source_text}
+</SOURCE_TEXT>
+
+<TRANSLATION>
+{translation_1}
+</TRANSLATION>
+
+When writing suggestions, pay attention to whether there are ways to improve the translation's \n\
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),\n\
+(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules, and ensuring there are no unnecessary repetitions),\n\
+(iii) style (by ensuring the translations reflect the style of the source text and takes into account any cultural context),\n\
+(iv) terminology (by ensuring terminology use is consistent and reflects the source text domain; and by only ensuring you use equivalent idioms {target_lang}).\n\
+
+Write a list of specific, helpful and constructive suggestions for improving the translation.
+Each suggestion should address one specific part of the translation.
+Output only the suggestions and nothing else."""
+
+    prompt = reflection_prompt.format(
+        source_lang=source_lang,
+        target_lang=target_lang,
+        source_text=source_text,
+        translation_1=translation_1,
+    )
+    reflection = get_completion(prompt, system_message=system_message)
+    return reflection
+
+
+def one_chunk_improve_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text: str,
+    translation_1: str,
+    reflection: str,
+) -> str:
+    """
+    Use the reflection to improve the translation, treating the entire text as one chunk.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for the translation.
+        source_text (str): The original text in the source language.
+        translation_1 (str): The initial translation of the source text.
+        reflection (str): Expert suggestions and constructive criticism for improving the translation.
+
+    Returns:
+        str: The improved translation based on the expert suggestions.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation editing from {source_lang} to {target_lang}."
+
+    prompt = f"""Your task is to carefully read, then edit, a translation from {source_lang} to {target_lang}, taking into
+account a list of expert suggestions and constructive criticisms.
+
+The source text, the initial translation, and the expert linguist suggestions are delimited by XML tags <SOURCE_TEXT></SOURCE_TEXT>, <TRANSLATION></TRANSLATION> and <EXPERT_SUGGESTIONS></EXPERT_SUGGESTIONS> \
+as follows:
+
+<SOURCE_TEXT>
+{source_text}
+</SOURCE_TEXT>
+
+<TRANSLATION>
+{translation_1}
+</TRANSLATION>
+
+<EXPERT_SUGGESTIONS>
+{reflection}
+</EXPERT_SUGGESTIONS>
+
+Please take into account the expert suggestions when editing the translation. Edit the translation by ensuring:
+
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),
+(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules and ensuring there are no unnecessary repetitions), \
+(iii) style (by ensuring the translations reflect the style of the source text)
+(iv) terminology (inappropriate for context, inconsistent use), or
+(v) other errors.
+
+Output only the new translation and nothing else."""
+
+    translation_2 = get_completion(prompt, system_message)
+
+    return translation_2
+
+
+def one_chunk_translate_text(
+    source_lang: str, target_lang: str, source_text: str, country: str = ""
+) -> str:
+    """
+    Translate a single chunk of text from the source language to the target language.
+
+    This function performs a two-step translation process:
+    1. Get an initial translation of the source text.
+    2. Reflect on the initial translation and generate an improved translation.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for the translation.
+        source_text (str): The text to be translated.
+        country (str): Country specified for target language.
+    Returns:
+        str: The improved translation of the source text.
+    """
+    translation_1 = one_chunk_initial_translation(
+        source_lang, target_lang, source_text
+    )
+
+    reflection = one_chunk_reflect_on_translation(
+        source_lang, target_lang, source_text, translation_1, country
+    )
+    translation_2 = one_chunk_improve_translation(
+        source_lang, target_lang, source_text, translation_1, reflection
+    )
+
+    return translation_2
+
+
+def num_tokens_in_string(
+    input_str: str, encoding_name: str = "cl100k_base"
+) -> int:
+    """
+    Calculate the number of tokens in a given string using a specified encoding.
+
+    Args:
+        str (str): The input string to be tokenized.
+        encoding_name (str, optional): The name of the encoding to use. Defaults to "cl100k_base",
+            which is the most commonly used encoder (used by GPT-4).
+
+    Returns:
+        int: The number of tokens in the input string.
+
+    Example:
+        >>> text = "Hello, how are you?"
+        >>> num_tokens = num_tokens_in_string(text)
+        >>> print(num_tokens)
+        5
+    """
+    encoding = tiktoken.get_encoding(encoding_name)
+    num_tokens = len(encoding.encode(input_str))
+    return num_tokens
+
+
+def multichunk_initial_translation(
+    source_lang: str, target_lang: str, source_text_chunks: List[str]
+) -> List[str]:
+    """
+    Translate a text in multiple chunks from the source language to the target language.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for translation.
+        source_text_chunks (List[str]): A list of text chunks to be translated.
+
+    Returns:
+        List[str]: A list of translated text chunks.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation from {source_lang} to {target_lang}."
+
+    translation_prompt = """Your task is provide a professional translation from {source_lang} to {target_lang} of PART of a text.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>. Translate only the part within the source text
+delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS>. You can use the rest of the source text as context, but do not translate any
+of the other text. Do not output anything other than the translation of the indicated part of the text.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, you should translate only this part of the text, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+Output only the translation of the portion you are asked to translate, and nothing else.
+"""
+
+    translation_chunks = []
+    for i in range(len(source_text_chunks)):
+        # Will translate chunk i
+        tagged_text = (
+            "".join(source_text_chunks[0:i])
+            + "<TRANSLATE_THIS>"
+            + source_text_chunks[i]
+            + "</TRANSLATE_THIS>"
+            + "".join(source_text_chunks[i + 1 :])
+        )
+
+        prompt = translation_prompt.format(
+            source_lang=source_lang,
+            target_lang=target_lang,
+            tagged_text=tagged_text,
+            chunk_to_translate=source_text_chunks[i],
+        )
+
+        translation = get_completion(prompt, system_message=system_message)
+        translation_chunks.append(translation)
+
+    return translation_chunks
+
+
+def multichunk_reflect_on_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text_chunks: List[str],
+    translation_1_chunks: List[str],
+    country: str = "",
+) -> List[str]:
+    """
+    Provides constructive criticism and suggestions for improving a partial translation.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language of the translation.
+        source_text_chunks (List[str]): The source text divided into chunks.
+        translation_1_chunks (List[str]): The translated chunks corresponding to the source text chunks.
+        country (str): Country specified for target language.
+
+    Returns:
+        List[str]: A list of reflections containing suggestions for improving each translated chunk.
+    """
+
+    system_message = f"You are an expert linguist specializing in translation from {source_lang} to {target_lang}. \
+You will be provided with a source text and its translation and your goal is to improve the translation."
+
+    if country != "":
+        reflection_prompt = """Your task is to carefully read a source text and part of a translation of that text from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions for improving the translation.
+The final style and tone of the translation should match the style of {target_lang} colloquially spoken in {country}.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>, and the part that has been translated
+is delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS> within the source text. You can use the rest of the source text
+as context for critiquing the translated part.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, only part of the text is being translated, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+The translation of the indicated part, delimited below by <TRANSLATION> and </TRANSLATION>, is as follows:
+<TRANSLATION>
+{translation_1_chunk}
+</TRANSLATION>
+
+When writing suggestions, pay attention to whether there are ways to improve the translation's:\n\
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),\n\
+(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules, and ensuring there are no unnecessary repetitions),\n\
+(iii) style (by ensuring the translations reflect the style of the source text and takes into account any cultural context),\n\
+(iv) terminology (by ensuring terminology use is consistent and reflects the source text domain; and by only ensuring you use equivalent idioms {target_lang}).\n\
+
+Write a list of specific, helpful and constructive suggestions for improving the translation.
+Each suggestion should address one specific part of the translation.
+Output only the suggestions and nothing else."""
+
+    else:
+        reflection_prompt = """Your task is to carefully read a source text and part of a translation of that text from {source_lang} to {target_lang}, and then give constructive criticism and helpful suggestions for improving the translation.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>, and the part that has been translated
+is delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS> within the source text. You can use the rest of the source text
+as context for critiquing the translated part.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, only part of the text is being translated, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+The translation of the indicated part, delimited below by <TRANSLATION> and </TRANSLATION>, is as follows:
+<TRANSLATION>
+{translation_1_chunk}
+</TRANSLATION>
+
+When writing suggestions, pay attention to whether there are ways to improve the translation's:\n\
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),\n\
+(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules, and ensuring there are no unnecessary repetitions),\n\
+(iii) style (by ensuring the translations reflect the style of the source text and takes into account any cultural context),\n\
+(iv) terminology (by ensuring terminology use is consistent and reflects the source text domain; and by only ensuring you use equivalent idioms {target_lang}).\n\
+
+Write a list of specific, helpful and constructive suggestions for improving the translation.
+Each suggestion should address one specific part of the translation.
+Output only the suggestions and nothing else."""
+
+    reflection_chunks = []
+    for i in range(len(source_text_chunks)):
+        # Will translate chunk i
+        tagged_text = (
+            "".join(source_text_chunks[0:i])
+            + "<TRANSLATE_THIS>"
+            + source_text_chunks[i]
+            + "</TRANSLATE_THIS>"
+            + "".join(source_text_chunks[i + 1 :])
+        )
+        if country != "":
+            prompt = reflection_prompt.format(
+                source_lang=source_lang,
+                target_lang=target_lang,
+                tagged_text=tagged_text,
+                chunk_to_translate=source_text_chunks[i],
+                translation_1_chunk=translation_1_chunks[i],
+                country=country,
+            )
+        else:
+            prompt = reflection_prompt.format(
+                source_lang=source_lang,
+                target_lang=target_lang,
+                tagged_text=tagged_text,
+                chunk_to_translate=source_text_chunks[i],
+                translation_1_chunk=translation_1_chunks[i],
+            )
+
+        reflection = get_completion(prompt, system_message=system_message)
+        reflection_chunks.append(reflection)
+
+    return reflection_chunks
+
+
+def multichunk_improve_translation(
+    source_lang: str,
+    target_lang: str,
+    source_text_chunks: List[str],
+    translation_1_chunks: List[str],
+    reflection_chunks: List[str],
+) -> List[str]:
+    """
+    Improves the translation of a text from source language to target language by considering expert suggestions.
+
+    Args:
+        source_lang (str): The source language of the text.
+        target_lang (str): The target language for translation.
+        source_text_chunks (List[str]): The source text divided into chunks.
+        translation_1_chunks (List[str]): The initial translation of each chunk.
+        reflection_chunks (List[str]): Expert suggestions for improving each translated chunk.
+
+    Returns:
+        List[str]: The improved translation of each chunk.
+    """
+
+    system_message = f"You are an expert linguist, specializing in translation editing from {source_lang} to {target_lang}."
+
+    improvement_prompt = """Your task is to carefully read, then improve, a translation from {source_lang} to {target_lang}, taking into
+account a set of expert suggestions and constructive critisms. Below, the source text, initial translation, and expert suggestions are provided.
+
+The source text is below, delimited by XML tags <SOURCE_TEXT> and </SOURCE_TEXT>, and the part that has been translated
+is delimited by <TRANSLATE_THIS> and </TRANSLATE_THIS> within the source text. You can use the rest of the source text
+as context, but need to provide a translation only of the part indicated by <TRANSLATE_THIS> and </TRANSLATE_THIS>.
+
+<SOURCE_TEXT>
+{tagged_text}
+</SOURCE_TEXT>
+
+To reiterate, only part of the text is being translated, shown here again between <TRANSLATE_THIS> and </TRANSLATE_THIS>:
+<TRANSLATE_THIS>
+{chunk_to_translate}
+</TRANSLATE_THIS>
+
+The translation of the indicated part, delimited below by <TRANSLATION> and </TRANSLATION>, is as follows:
+<TRANSLATION>
+{translation_1_chunk}
+</TRANSLATION>
+
+The expert translations of the indicated part, delimited below by <EXPERT_SUGGESTIONS> and </EXPERT_SUGGESTIONS>, is as follows:
+<EXPERT_SUGGESTIONS>
+{reflection_chunk}
+</EXPERT_SUGGESTIONS>
+
+Taking into account the expert suggestions rewrite the translation to improve it, paying attention
+to whether there are ways to improve the translation's
+
+(i) accuracy (by correcting errors of addition, mistranslation, omission, or untranslated text),
+(ii) fluency (by applying {target_lang} grammar, spelling and punctuation rules and ensuring there are no unnecessary repetitions), \
+(iii) style (by ensuring the translations reflect the style of the source text)
+(iv) terminology (inappropriate for context, inconsistent use), or
+(v) other errors.
+
+Output only the new translation of the indicated part and nothing else."""
+
+    translation_2_chunks = []
+    for i in range(len(source_text_chunks)):
+        # Will translate chunk i
+        tagged_text = (
+            "".join(source_text_chunks[0:i])
+            + "<TRANSLATE_THIS>"
+            + source_text_chunks[i]
+            + "</TRANSLATE_THIS>"
+            + "".join(source_text_chunks[i + 1 :])
+        )
+
+        prompt = improvement_prompt.format(
+            source_lang=source_lang,
+            target_lang=target_lang,
+            tagged_text=tagged_text,
+            chunk_to_translate=source_text_chunks[i],
+            translation_1_chunk=translation_1_chunks[i],
+            reflection_chunk=reflection_chunks[i],
+        )
+
+        translation_2 = get_completion(prompt, system_message=system_message)
+        translation_2_chunks.append(translation_2)
+
+    return translation_2_chunks
+
+
+def multichunk_translation(
+    source_lang, target_lang, source_text_chunks, country: str = ""
+):
+    """
+    Improves the translation of multiple text chunks based on the initial translation and reflection.
+
+    Args:
+        source_lang (str): The source language of the text chunks.
+        target_lang (str): The target language for translation.
+        source_text_chunks (List[str]): The list of source text chunks to be translated.
+        translation_1_chunks (List[str]): The list of initial translations for each source text chunk.
+        reflection_chunks (List[str]): The list of reflections on the initial translations.
+        country (str): Country specified for target language
+    Returns:
+        List[str]: The list of improved translations for each source text chunk.
+    """
+
+    translation_1_chunks = multichunk_initial_translation(
+        source_lang, target_lang, source_text_chunks
+    )
+
+    reflection_chunks = multichunk_reflect_on_translation(
+        source_lang,
+        target_lang,
+        source_text_chunks,
+        translation_1_chunks,
+        country,
+    )
+
+    translation_2_chunks = multichunk_improve_translation(
+        source_lang,
+        target_lang,
+        source_text_chunks,
+        translation_1_chunks,
+        reflection_chunks,
+    )
+
+    return translation_2_chunks
+
+
+def calculate_chunk_size(token_count: int, token_limit: int) -> int:
+    """
+    Calculate the chunk size based on the token count and token limit.
+
+    Args:
+        token_count (int): The total number of tokens.
+        token_limit (int): The maximum number of tokens allowed per chunk.
+
+    Returns:
+        int: The calculated chunk size.
+
+    Description:
+        This function calculates the chunk size based on the given token count and token limit.
+        If the token count is less than or equal to the token limit, the function returns the token count as the chunk size.
+        Otherwise, it calculates the number of chunks needed to accommodate all the tokens within the token limit.
+        The chunk size is determined by dividing the token limit by the number of chunks.
+        If there are remaining tokens after dividing the token count by the token limit,
+        the chunk size is adjusted by adding the remaining tokens divided by the number of chunks.
+
+    Example:
+        >>> calculate_chunk_size(1000, 500)
+        500
+        >>> calculate_chunk_size(1530, 500)
+        389
+        >>> calculate_chunk_size(2242, 500)
+        496
+    """
+
+    if token_count <= token_limit:
+        return token_count
+
+    num_chunks = (token_count + token_limit - 1) // token_limit
+    chunk_size = token_count // num_chunks
+
+    remaining_tokens = token_count % token_limit
+    if remaining_tokens > 0:
+        chunk_size += remaining_tokens // num_chunks
+
+    return chunk_size
+
+
+def translate(
+    source_lang,
+    target_lang,
+    source_text,
+    country,
+    max_tokens=MAX_TOKENS_PER_CHUNK,
+):
+    """Translate the source_text from source_lang to target_lang."""
+
+    num_tokens_in_text = num_tokens_in_string(source_text)
+
+    ic(num_tokens_in_text)
+
+    if num_tokens_in_text < max_tokens:
+        ic("Translating text as single chunk")
+
+        final_translation = one_chunk_translate_text(
+            source_lang, target_lang, source_text, country
+        )
+
+        return final_translation
+
+    else:
+        ic("Translating text as multiple chunks")
+
+        token_size = calculate_chunk_size(
+            token_count=num_tokens_in_text, token_limit=max_tokens
+        )
+
+        ic(token_size)
+
+        text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
+            model_name="gpt-4",
+            chunk_size=token_size,
+            chunk_overlap=0,
+        )
+
+        source_text_chunks = text_splitter.split_text(source_text)
+
+        translation_2_chunks = multichunk_translation(
+            source_lang, target_lang, source_text_chunks, country
+        )
+
+        return "".join(translation_2_chunks)