Views
No views yet
1from transformers import (
2 AutoModelForCausalLM,
3 AutoTokenizer,
4 BitsAndBytesConfig,
5)
6from peft import PeftModel
7import torch
8from tqdm import tqdm
9import json
10
11HF_TOKEN = "Hugging Face Token"
12model_id = "llm-jp/llm-jp-3-13b"
13base_model_id = "llm-jp/llm-jp-3-13b"
14adapter_id = "A-Haya/llm-jp-3-13b-finetune"
15
16# QLoRA config
17bnb_config = BitsAndBytesConfig(
18 load_in_4bit=True,
19 bnb_4bit_quant_type="nf4",
20 bnb_4bit_compute_dtype=torch.bfloat16,
21)
22
23# Load model
24model = AutoModelForCausalLM.from_pretrained(
25 model_id,
26 quantization_config=bnb_config,
27 device_map="auto",
28 token = HF_TOKEN
29)
30
31# Load tokenizer
32tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token = HF_TOKEN)
33
34model = PeftModel.from_pretrained(model, adapter_id, token = HF_TOKEN)
35
36
37datasets = []
38with open("your_datasets.jsonl", "r") as f:
39 item = ""
40 for line in f:
41 line = line.strip()
42 item += line
43 if item.endswith("}"):
44 datasets.append(json.loads(item))
45 item = ""
46
47# 推論の実行
48results = []
49for data in tqdm(datasets):
50
51 input = data["input"]
52
53 prompt = f"""### 指示
54 {input}
55 ### 回答
56 """
57
58 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
59 attention_mask = torch.ones_like(tokenized_input)
60 with torch.no_grad():
61 outputs = model.generate(
62 tokenized_input,
63 attention_mask=attention_mask,
64 max_new_tokens=100,
65 do_sample=False,
66 repetition_penalty=1.2,
67 pad_token_id=tokenizer.eos_token_id
68 )[0]
69 output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)
70
71 results.append({"task_id": data["task_id"], "input": input, "output": output})