Views
No views yet
1
2from transformers import (
3 AutoModelForCausalLM,
4 AutoTokenizer,
5 BitsAndBytesConfig,
6)
7from peft import PeftModel
8import torch
9from tqdm import tqdm
10import json
11
12HF_TOKEN = # your token
13
14model_id = "llm-jp/llm-jp-3-13b"
15adapter_id = "Wangmio/llm-jpllm-jp-3-13b-f2-Wmq-2"
16adapter_dpo_id = "Wangmio/llm-jpllm-jp-3-13b-f2-Wmq-2-dpo"
17
18bnb_config = BitsAndBytesConfig(
19 load_in_4bit=True,
20 bnb_4bit_quant_type="nf4",
21 bnb_4bit_compute_dtype=torch.bfloat16,
22)
23
24# Load model
25model = AutoModelForCausalLM.from_pretrained(
26 model_id,
27 quantization_config=bnb_config,
28 device_map="auto",
29 token = HF_TOKEN
30)
31
32# Load tokenizer
33tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, token = HF_TOKEN)
34
35model = PeftModel.from_pretrained(model, adapter_dpo_id, token = HF_TOKEN)
36
37datasets = []
38with open("./elyza-tasks-100-TV_0.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
47results = []
48for data in tqdm(datasets):
49
50 input = data['input']
51
52 prompt = f"""### 指示
53 {input}
54 ### 回答
55 """
56
57 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
58 attention_mask = torch.ones_like(tokenized_input)
59 with torch.no_grad():
60 outputs = model.generate(
61 tokenized_input,
62 attention_mask=attention_mask,
63 max_new_tokens=100,
64 do_sample=False,
65 repetition_penalty=1.2,
66 pad_token_id=tokenizer.eos_token_id
67 )[0]
68 output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)
69
70 results.append({"task_id": data["task_id"], "output": output})
71