Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2from transformers.generation import GenerationConfig
3
4model_path = 'SakuraLLM/LN-Korean-14B-v0.2.1'
5tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
6model = AutoModelForCausalLM.from_pretrained(model_path, device_map='auto', trust_remote_code=True).eval()
7model.generation_config = GenerationConfig.from_pretrained(model_path, trust_remote_code=True)
8
9# 段落之间用\n分隔
10text = '''여자애들이 자신들의 첫 경험에 대한 이야기를 하는 걸 들은 적이 있는가.
11물론 여기서 첫 경험이라는 것은 처음으로 야자를 쨌다든가 처음으로 술을 마셔 봤다든가 그런 것이 아니라, 명실공히 그렇고 그런 의미에서의 첫 경험이다.
12“우, 우리가…… 처음으로 그, 그걸 한 거는 말이야.”
13그렇게 말한 것은 소파에 앉아 있는 갈색 교복의 소녀였다. 둥근 얼굴에 커다란 갈색 눈동자를 지닌, 부드러운 머리카락을 어깨 위로 늘어뜨리고 있는 소녀다. 전반적으로 얌전한 모범생 같아 보이는 인상이고 몸집도 아담한 편이지만, 교복 상의를 매혹적으로 부풀어 오르게 하고 있는 가슴만큼은 얌전하지도 아담하지도 않았다. 몸을 움츠린 자세 탓에 두 팔이 가슴을 양옆에서 압박하고 있어, 몸을 움직일 때마다 그 윤곽이 부드럽게 일그러졌다.'''
14
15# 文本长度控制在1024以内
16assert len(text) < 1024
17
18messages = [
19 {'role': 'user', 'content': f'翻译成中文:\n{text}'}
20]
21
22text = tokenizer.apply_chat_template(
23 messages,
24 tokenize=False,
25 add_generation_prompt=True
26)
27
28model_inputs = tokenizer([text], return_tensors='pt').to('cuda')
29
30generated_ids = model.generate(
31 model_inputs.input_ids,
32 max_new_tokens=1024
33)
34
35generated_ids = [
36 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
37]
38
39response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
40print(response)