Views
No views yet
1"""
2<s>以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい
3
4### 指示:
5{instruction}
6
7### 応答:
8"""1import torch
2from transformers import (
3 AutoTokenizer,
4 AutoModelForCausalLM,
5 BitsAndBytesConfig,
6)
7from peft import LoraConfig, PeftModel
8from transformers import TextStreamer
9
10
11BASE_MODEL = "llm-jp/llm-jp-3-13b"
12PEFT_MODEL = "togepi55/llm-jp-3-13b-it"
13
14tokenizer = AutoTokenizer.from_pretrained(PEFT_MODEL)
15bnb_config = BitsAndBytesConfig(
16 load_in_4bit=True,
17 bnb_4bit_compute_dtype=torch.float16,
18 bnb_4bit_quant_type="nf4",
19 bnb_4bit_use_double_quant=False,
20)
21
22base_model = AutoModelForCausalLM.from_pretrained(
23 BASE_MODEL,
24 device_map="auto",
25 quantization_config=bnb_config,
26 torch_dtype="auto",
27 trust_remote_code=True,
28 )
29
30
31model = PeftModel.from_pretrained(base_model, PEFT_MODEL)
32
33streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
34
35
36instruction = "東京は日本の"
37
38prompt = f"<s>以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい\n\n### 指示:\n{instruction}\n\n### 応答:\n"
39print(prompt)
40model_input = tokenizer(prompt, return_tensors="pt").to(model.device)
41input_ids = model_input["input_ids"]
42
43model.eval()
44with torch.no_grad():
45 result = model.generate(
46 input_ids,
47 max_new_tokens=300,
48 attention_mask = model_input.attention_mask,
49 pad_token_id=tokenizer.pad_token_id,
50 eos_token_id=tokenizer.eos_token_id,
51 do_sample=False,
52 streamer=streamer,
53 repetition_penalty=1.02,
54 )
55 print("----"*20)
56 del input_ids
57 torch.cuda.empty_cache()