Views
No views yet

{BOS}[TRANS]\n{source_sentence}\n[/TRANS]\n[{target_language}]\nsource_sentence: The sentence you want to translate.target_language: The target language you want to translate to. Use "ZH" for Traditional Chinese, "EN" for English, "POJ" for Taiwanese Hokkien POJ, "HL" for Taiwanese Hokkien Hanlo, and "HAN" for Taiwanese Hokkien Hanzi.1from transformers import AutoModelForCausalLM, AutoTokenizer, TextGenerationPipeline
2import torch
3import accelerate
4
5def get_pipeline(path:str, tokenizer:AutoTokenizer, accelerator:accelerate.Accelerator) -> TextGenerationPipeline:
6 model = AutoModelForCausalLM.from_pretrained(
7 path, torch_dtype=torch.float16, device_map='auto', trust_remote_code=True)
8
9 terminators = [tokenizer.eos_token_id, tokenizer.pad_token_id]
10
11 pipeline = TextGenerationPipeline(model = model, tokenizer = tokenizer, num_workers=accelerator.state.num_processes*4, pad_token_id=tokenizer.pad_token_id, eos_token_id=terminators)
12
13 return pipeline
14
15model_dir = "Bohanlu/Taigi-Llama-2-Translator-7B" # or "Bohanlu/Taigi-Llama-2-Translator-13B" for the 13B model
16tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=False)
17
18accelerator = accelerate.Accelerator()
19pipe = get_pipeline(model_dir, tokenizer, accelerator)
20
21PROMPT_TEMPLATE = "[TRANS]\n{source_sentence}\n[/TRANS]\n[{target_language}]\n"
22
23def translate(source_sentence:str, target_language:str) -> str:
24 prompt = PROMPT_TEMPLATE.format(source_sentence=source_sentence, target_language=target_language)
25 out = pipe(prompt, return_full_text=False, repetition_penalty=1.1, do_sample=False)[0]['generated_text']
26 return out[:out.find("[/")].strip()
27
28source_sentence = "How are you today?"
29
30print("To Hanzi: " + translate(source_sentence, "HAN"))
31# Output: To Hanzi: 你今仔日好無?
32
33print("To POJ: " + translate(source_sentence, "POJ"))
34# Output: To POJ: Lí kin-á-ji̍t án-chóaⁿ?
35
36print("To Traditional Chinese: " + translate(source_sentence, "ZH"))
37# Output: To Traditional Chinese: 你今天好嗎?
38
39print("To Hanlo: " + translate(source_sentence, "HL"))
40# Output: To Hanlo: 你今仔日好無?@misc{lu2024enhancing,
title={Enhancing Taiwanese Hokkien Dual Translation by Exploring and Standardizing of Four Writing Systems},
author={Bo-Han Lu and Yi-Hsuan Lin and En-Shiun Annie Lee and Richard Tzong-Han Tsai},
year={2024},
eprint={2403.12024},
archivePrefix={arXiv},
primaryClass={cs.CL}
}