File size: 5,965 Bytes
3ea78af 6c3e0fe 3ea78af |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from argparse import ArgumentParser
import random
import time
import re
from tqdm import tqdm
def get_novel_text_list(data_path, text_length):
data_list = list()
with open(data_path, 'r', encoding="utf-8") as f:
data = f.read()
data = data.replace(" ", "")
data_raw = re.sub('\n+', '\n', data)
data = data_raw.strip().split("\n")
i = 0
while i < len(data):
r = random.randint(int(text_length/2), text_length)
text = ""
while len(text) < r:
if i >= len(data):
break
if len(text) > max(- len(data[i]) + r, 0):
break
else:
text += data[i] + "\n"
i += 1
text = text.strip()
data_list.append(text)
return data_raw, data_list
def get_prompt(input, model_version):
if model_version == '0.5' or model_version == '0.8':
prompt = "<reserved_106>将下面的日文文本翻译成中文：" + input + "<reserved_107>"
return prompt
if model_version == '0.7':
prompt = f"<|im_start|>user\n将下面的日文文本翻译成中文：{input}<|im_end|>\n<|im_start|>assistant\n"
return prompt
if model_version == '0.1':
prompt = "Human: \n将下面的日文文本翻译成中文：" + input + "\n\nAssistant: \n"
return prompt
if model_version == '0.4':
prompt = "User: 将下面的日文文本翻译成中文：" + input + "\nAssistant: "
return prompt
raise ValueError(f"Wrong model version{model_version}, please view https://huggingface.co/sakuraumi/Sakura-13B-Galgame")
def split_response(response, model_version):
response = response.replace("</s>", "")
if model_version == '0.5' or model_version == '0.8':
output = response.split("<reserved_107>")[1]
return output
if model_version == '0.7':
output = response.split("<|im_start|>assistant\n")[1]
return output
if model_version == '0.1':
output = response.split("\n\nAssistant: \n")[1]
return output
if model_version == '0.4':
output = response.split("\nAssistant: ")[1]
return output
raise ValueError(f"Wrong model version{model_version}, please view https://huggingface.co/sakuraumi/Sakura-13B-Galgame")
def get_model_response(model: AutoModelForCausalLM, tokenizer: AutoTokenizer, prompt: str, model_version: str, generation_config: GenerationConfig):
generation = model.generate(**tokenizer(prompt, return_tensors="pt").to(model.device), generation_config=generation_config)[0]
response = tokenizer.decode(generation)
output = split_response(response, model_version)
return output
def get_compare_text(source_text, translated_text):
source_text_list = source_text.strip().split("\n")
translated_text_list = translated_text.strip().split("\n")
output_text = ""
if len(source_text_list) != len(translated_text_list):
print("error occurred when output compared text, fallback to output only translated text.")
return translated_text
else:
for i in range(len(source_text_list)):
output_text += source_text_list[i] + "\n" + translated_text_list[i] + "\n\n"
output_text = output_text.strip()
return output_text
def main():
parser = ArgumentParser()
parser.add_argument("--model_name_or_path", type=str, default="SakuraLLM/Sakura-13B-LNovel-v0.8", help="model huggingface id or local path.")
parser.add_argument("--use_gptq_model", action="store_true", help="whether your model is gptq quantized.")
parser.add_argument("--model_version", type=str, default="0.8", help="model version written on huggingface readme, now we have ['0.1', '0.4', '0.5', '0.7', '0.8']")
parser.add_argument("--data_path", type=str, default="data.txt", help="file path of the text you want to translate.")
parser.add_argument("--output_path", type=str, default="data_translated.txt", help="save path of the text model translated.")
parser.add_argument("--text_length", type=int, default=512, help="input max length in each inference.")
parser.add_argument("--compare_text", action="store_true", help="whether to output with both source text and translated text in order to compare.")
args = parser.parse_args()
if args.use_gptq_model:
from auto_gptq import AutoGPTQForCausalLM
generation_config = GenerationConfig(
temperature=0.1,
top_p=0.3,
top_k=40,
num_beams=1,
bos_token_id=1,
eos_token_id=2,
pad_token_id=0,
max_new_tokens=1024,
min_new_tokens=1,
do_sample=True
)
print("loading...")
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, use_fast=False, trust_remote_code=True)
if args.use_gptq_model:
model = AutoGPTQForCausalLM.from_quantized(args.model_name_or_path, device="cuda:0", trust_remote_code=True)
else:
model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, device_map="auto", trust_remote_code=True)
print("translating...")
start = time.time()
data_raw, data_list = get_novel_text_list(args.data_path, args.text_length)
data = ""
for d in tqdm(data_list):
prompt = get_prompt(d, args.model_version)
output = get_model_response(model, tokenizer, prompt, args.model_version, generation_config)
data += output.strip() + "\n"
end = time.time()
print("translation completed, used time: ", end-start)
print("saving...")
if args.compare_text:
with open(args.output_path, 'w', encoding='utf-8') as f_w:
f_w.write(get_compare_text(data_raw, data))
else:
with open(args.output_path, 'w', encoding='utf-8') as f_w:
f_w.write(data)
print("completed.")
if __name__ == "__main__":
main() |