Views
No views yet
1import os, torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
3from peft import PeftModel
4
5HF_TOKEN = os.environ.get('HF_LLM_API_KEY')
6
7base_model_id = "llm-jp/llm-jp-3-13b"
8adapter_id = "5h11Vj0/llm-jp-3-13b_ft_nf4"
9
10bnb_config = BitsAndBytesConfig(
11 load_in_4bit=True,
12 bnb_4bit_quant_type="nf4",
13 bnb_4bit_compute_dtype=torch.bfloat16
14)
15model = AutoModelForCausalLM.from_pretrained(
16 base_model_id,
17 quantization_config=bnb_config,
18 device_map="auto",
19 token = HF_TOKEN
20)
21model = PeftModel.from_pretrained(model, adapter_id, token=HF_TOKEN)
22tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True, token=HF_TOKEN)
23
24input_text = "連濁とはなんですか?"
25prompt = f"""### 指示
26{input_text}
27### 回答
28"""
29
30tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
31attention_mask = torch.ones_like(tokenized_input)
32
33with torch.no_grad():
34 outputs = model.generate(
35 tokenized_input,
36 attention_mask=attention_mask,
37 max_new_tokens=100,
38 do_sample=False,
39 repetition_penalty=1.2,
40 pad_token_id=tokenizer.eos_token_id
41 )[0]
42output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)
43print(output)