From bfb6b93d16f133061eaf06e98d979c93a93d26b0 Mon Sep 17 00:00:00 2001 From: csh <458761603@qq.com> Date: Mon, 24 Jun 2024 21:56:48 +0800 Subject: [PATCH 1/3] Allow change model Signed-off-by: csh <458761603@qq.com> --- src/translation_agent/utils.py | 53 ++++++++++++++++++++++++---------- 1 file changed, 38 insertions(+), 15 deletions(-) diff --git a/src/translation_agent/utils.py b/src/translation_agent/utils.py index 1256281..51c2869 100755 --- a/src/translation_agent/utils.py +++ b/src/translation_agent/utils.py @@ -71,7 +71,7 @@ def get_completion( def one_chunk_initial_translation( - source_lang: str, target_lang: str, source_text: str + source_lang: str, target_lang: str, source_text: str, model: str = "gpt-4-turbo", ) -> str: """ Translate the entire text as one chunk using an LLM. @@ -80,6 +80,8 @@ def one_chunk_initial_translation( 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. + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". Returns: str: The translated text. @@ -95,7 +97,7 @@ def one_chunk_initial_translation( prompt = translation_prompt.format(source_text=source_text) - translation = get_completion(prompt, system_message=system_message) + translation = get_completion(prompt, system_message=system_message, model=model) return translation @@ -106,6 +108,7 @@ def one_chunk_reflect_on_translation( source_text: str, translation_1: str, country: str = "", + model: str = "gpt-4-turbo", ) -> str: """ Use an LLM to reflect on the translation, treating the entire text as one chunk. @@ -116,6 +119,8 @@ def one_chunk_reflect_on_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. + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". Returns: str: The LLM's reflection on the translation, providing constructive criticism and suggestions for improvement. @@ -177,7 +182,7 @@ def one_chunk_reflect_on_translation( source_text=source_text, translation_1=translation_1, ) - reflection = get_completion(prompt, system_message=system_message) + reflection = get_completion(prompt, system_message=system_message, model=model) return reflection @@ -187,6 +192,7 @@ def one_chunk_improve_translation( source_text: str, translation_1: str, reflection: str, + model: str = "gpt-4-turbo", ) -> str: """ Use the reflection to improve the translation, treating the entire text as one chunk. @@ -197,6 +203,8 @@ def one_chunk_improve_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. + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". Returns: str: The improved translation based on the expert suggestions. @@ -232,13 +240,13 @@ def one_chunk_improve_translation( Output only the new translation and nothing else.""" - translation_2 = get_completion(prompt, system_message) + translation_2 = get_completion(prompt, system_message, model=model) return translation_2 def one_chunk_translate_text( - source_lang: str, target_lang: str, source_text: str, country: str = "" + source_lang: str, target_lang: str, source_text: str, country: str = "", model: str = "gpt-4-turbo", ) -> str: """ Translate a single chunk of text from the source language to the target language. @@ -252,18 +260,20 @@ def one_chunk_translate_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. + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". Returns: str: The improved translation of the source text. """ translation_1 = one_chunk_initial_translation( - source_lang, target_lang, source_text + source_lang, target_lang, source_text, model ) reflection = one_chunk_reflect_on_translation( - source_lang, target_lang, source_text, translation_1, country + source_lang, target_lang, source_text, translation_1, country, model ) translation_2 = one_chunk_improve_translation( - source_lang, target_lang, source_text, translation_1, reflection + source_lang, target_lang, source_text, translation_1, reflection, model ) return translation_2 @@ -295,7 +305,7 @@ def num_tokens_in_string( def multichunk_initial_translation( - source_lang: str, target_lang: str, source_text_chunks: List[str] + source_lang: str, target_lang: str, source_text_chunks: List[str], model: str = "gpt-4-turbo", ) -> List[str]: """ Translate a text in multiple chunks from the source language to the target language. @@ -304,6 +314,8 @@ def multichunk_initial_translation( 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. + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". Returns: List[str]: A list of translated text chunks. @@ -347,7 +359,7 @@ def multichunk_initial_translation( chunk_to_translate=source_text_chunks[i], ) - translation = get_completion(prompt, system_message=system_message) + translation = get_completion(prompt, system_message=system_message, model=model) translation_chunks.append(translation) return translation_chunks @@ -359,6 +371,7 @@ def multichunk_reflect_on_translation( source_text_chunks: List[str], translation_1_chunks: List[str], country: str = "", + model: str = "gpt-4-turbo", ) -> List[str]: """ Provides constructive criticism and suggestions for improving a partial translation. @@ -369,6 +382,8 @@ def multichunk_reflect_on_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. + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". Returns: List[str]: A list of reflections containing suggestions for improving each translated chunk. @@ -468,7 +483,7 @@ def multichunk_reflect_on_translation( translation_1_chunk=translation_1_chunks[i], ) - reflection = get_completion(prompt, system_message=system_message) + reflection = get_completion(prompt, system_message=system_message, model=model) reflection_chunks.append(reflection) return reflection_chunks @@ -480,6 +495,7 @@ def multichunk_improve_translation( source_text_chunks: List[str], translation_1_chunks: List[str], reflection_chunks: List[str], + model: str = "gpt-4-turbo", ) -> List[str]: """ Improves the translation of a text from source language to target language by considering expert suggestions. @@ -490,6 +506,8 @@ def multichunk_improve_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. + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". Returns: List[str]: The improved translation of each chunk. @@ -554,14 +572,14 @@ def multichunk_improve_translation( reflection_chunk=reflection_chunks[i], ) - translation_2 = get_completion(prompt, system_message=system_message) + translation_2 = get_completion(prompt, system_message=system_message, model=model) translation_2_chunks.append(translation_2) return translation_2_chunks def multichunk_translation( - source_lang, target_lang, source_text_chunks, country: str = "" + source_lang, target_lang, source_text_chunks, country: str = "", model: str = "gpt-4-turbo", ): """ Improves the translation of multiple text chunks based on the initial translation and reflection. @@ -573,12 +591,14 @@ def multichunk_translation( 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 + model (str, optional): The name of the OpenAI model to use for generating the completion. + Defaults to "gpt-4-turbo". 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 + source_lang, target_lang, source_text_chunks, model ) reflection_chunks = multichunk_reflect_on_translation( @@ -587,6 +607,7 @@ def multichunk_translation( source_text_chunks, translation_1_chunks, country, + model, ) translation_2_chunks = multichunk_improve_translation( @@ -595,6 +616,7 @@ def multichunk_translation( source_text_chunks, translation_1_chunks, reflection_chunks, + model, ) return translation_2_chunks @@ -647,6 +669,7 @@ def translate( source_text, country, max_tokens=MAX_TOKENS_PER_CHUNK, + model="gpt-4-turbo", ): """Translate the source_text from source_lang to target_lang.""" @@ -658,7 +681,7 @@ def translate( ic("Translating text as single chunk") final_translation = one_chunk_translate_text( - source_lang, target_lang, source_text, country + source_lang, target_lang, source_text, country, model ) return final_translation From e2993b378d47b9240c4ae08e093e964a5e9145fd Mon Sep 17 00:00:00 2001 From: csh <458761603@qq.com> Date: Mon, 24 Jun 2024 22:26:34 +0800 Subject: [PATCH 2/3] Update README.md Signed-off-by: csh <458761603@qq.com> --- README.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/README.md b/README.md index 7f4db9f..a334ad1 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,21 @@ translation = ta.translate(source_lang, target_lang, source_text, country) ``` See examples/example_script.py for an example script to try out. +### Using OpenAI-compatible AI as backend: + +- Add `OPENAI_BASE_URL` to the environment variables or .env file to switch the API URL. +```bash +# export OPENAI_BASE_URL="https://api.openai.com/v1" +export OPENAI_BASE_URL="http://localhost:8080/v1 +``` + +- Select a model. +```python +import translation_agent as ta +source_lang, target_lang, country = "English", "Spanish", "Mexico" +translation = ta.translate(source_lang, target_lang, source_text, country, model = "llama-3-8b") +``` + ## License Translation Agent is released under the **MIT License**. You are free to use, modify, and distribute the code From 40a480f5198813e0ca40735e5fb8e60ff8f1f048 Mon Sep 17 00:00:00 2001 From: csh <458761603@qq.com> Date: Sat, 29 Jun 2024 11:15:31 +0800 Subject: [PATCH 3/3] Fix `ulits.translate`, `multichunk_translation` use model Signed-off-by: csh <458761603@qq.com> --- src/translation_agent/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/translation_agent/utils.py b/src/translation_agent/utils.py index 51c2869..5ed460f 100755 --- a/src/translation_agent/utils.py +++ b/src/translation_agent/utils.py @@ -704,7 +704,7 @@ def translate( source_text_chunks = text_splitter.split_text(source_text) translation_2_chunks = multichunk_translation( - source_lang, target_lang, source_text_chunks, country + source_lang, target_lang, source_text_chunks, country, model ) return "".join(translation_2_chunks)