Views
No views yet
1from vllm import LLM, SamplingParams
2from transformers import AutoTokenizer, pipeline
3
4BASE_MODEL = "sh2orc/gemma-1.1-korean-7b-it"
5
6llm = LLM(model=BASE_MODEL)
7
8tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
9tokenizer.pad_token = tokenizer.eos_token
10tokenizer.padding_side = 'right'
11
12instruction = '독도에 대해서 알려줘'
13
14messages = [
15 {
16 "role": "user",
17 "content": instruction
18 },
19]
20
21
22prompt_message = tokenizer.apply_chat_template(
23 messages,
24 tokenize=False,
25 add_generation_prompt=True,
26)
27
28eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<end_of_turn>")]
29
30outputs = llm.generate(prompt_message, SamplingParams(stop_token_ids=eos_token_id, temperature=0.2, top_p=0.8,max_tokens=4096))
31
32for output in outputs:
33 propt = output.prompt
34 generated_text = output.outputs[0].text
35 print(generated_text)
36