Views
No views yet
1import torch
2from typing import Literal
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5model_path = 'Mxode/NanoTranslator-immersive_translate-365M'
6
7model = AutoModelForCausalLM.from_pretrained(model_path).to('cuda:0', torch.bfloat16)
8tokenizer = AutoTokenizer.from_pretrained(model_path)
9
10def translate(
11 text: str,
12 to: Literal["chinese", "english"] = "chinese",
13 **kwargs
14):
15 generation_args = dict(
16 max_new_tokens = kwargs.pop("max_new_tokens", 512),
17 do_sample = kwargs.pop("do_sample", True),
18 temperature = kwargs.pop("temperature", 0.35),
19 top_p = kwargs.pop("top_p", 0.8),
20 top_k = kwargs.pop("top_k", 40),
21 **kwargs
22 )
23
24 prompt = """Translate the following source text to {to}. Output translation directly without any additional text.
25 Source Text: {text}
26
27 Translated Text:"""
28
29 messages = [
30 {"role": "system", "content": "You are a professional, authentic machine translation engine."},
31 {"role": "user", "content": prompt.format(to=to, text=text)}
32 ]
33 inputs = tokenizer.apply_chat_template(
34 messages,
35 tokenize=False,
36 add_generation_prompt=True
37 )
38 model_inputs = tokenizer([inputs], return_tensors="pt").to(model.device)
39
40 generated_ids = model.generate(model_inputs.input_ids, **generation_args)
41 generated_ids = [
42 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
43 ]
44
45 response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
46 return response
47
48text = "After a long day at work, I love to unwind by cooking a nice dinner and watching my favorite TV series. It really helps me relax and recharge for the next day."
49response = translate(text=text, to='chinese')
50print(f'Translation: {response}')
51
52"""
53Translation: 工作了一天，我喜欢吃一顿美味的晚餐，看我最喜欢的电视剧，这样做有助于我放松，补充能量。
54"""