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
12
13import os
14os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
15
16
17
18# hugging face tokenの設定
19HF_TOKEN = "your_token"
20
21# hugging faceへログイン
22from huggingface_hub import login
23login(token=HF_TOKEN)
24
25# 使用するモデルの設定
26model_id = "llm-jp/llm-jp-3-13b"
27adapter_id = "awaawattt/llm-jp-3-13b-finetune-all-data-reasoning-v15"
28
29
30# 量子化の設定
31bnb_config = BitsAndBytesConfig(
32 load_in_8bit=True,
33 llm_int8_threshold=6.0,
34 bnb_8bit_quant_type="nf4",
35 bnb_8bit_compute_dtype=torch.bfloat16,
36 bnb_8bit_use_double_quant=False,
37)
38
39
40# モデルとトークナイザーの読み込み
41model = AutoModelForCausalLM.from_pretrained(
42 model_id,
43 quantization_config=bnb_config,
44 device_map="auto",
45 torch_dtype=torch.bfloat16,
46 use_flash_attention_2=True, # FlashAttention-2を有効化
47 attn_implementation="flash_attention_2",
48)
49
50model = PeftModel.from_pretrained(model, adapter_id, trust_remote_code=True)
51
52tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
53
54
55# タスクとなるデータの読み込み。
56# omnicampusの開発環境では、左にタスクのjsonlをドラッグアンドドロップしてから実行。
57import json
58datasets = []
59with open("./elyza-tasks-100-TV_0.jsonl", "r") as f:
60 item = ""
61 for line in f:
62 line = line.strip()
63 item += line
64 if item.endswith("}"):
65 datasets.append(json.loads(item))
66 item = ""
67
68
69
70
71# モデルによるタスクの推論。
72from tqdm import tqdm
73results = []
74for data in tqdm(datasets):
75 input = data["input"]
76 prompt = f"""### 指示\n\n\n{input}\n\n### 回答\n"""
77 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
78 attention_mask = torch.ones_like(tokenized_input)
79 with torch.no_grad():
80 outputs = model.generate(
81 tokenized_input,
82 attention_mask=attention_mask,
83 max_new_tokens=4096,
84 do_sample=False,
85 repetition_penalty=1.2,
86 pad_token_id=tokenizer.eos_token_id,
87 use_cache=True,
88 )[0]
89 output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=False)
90 if ('<think>' in output) and ('</think>' in output):
91 output = output.split('</think>')[1]
92 reasoning = output.split('</think>')[0]
93 reasoning = reasoning.replace('<think>', '')
94 results.append({"task_id": data["task_id"], "input": input, "output": output, "reasoning": output})
95 # break
96 else:
97 if '<think>' in output:
98 output = output.replace('<think>', '')
99 print('think was not found at:', data["task_id"])
100 # print(output)
101 results.append({"task_id": data["task_id"], "input": input, "output": output, "reasoning": ''})
102
103
104# 回答ファイルの作成
105import re
106jsonl_id = re.sub(".*/", "", adapter_id)
107with open(f"./{jsonl_id}-outputs.jsonl", 'w', encoding='utf-8') as f:
108 for result in results:
109 json.dump(result, f, ensure_ascii=False) # ensure_ascii=False for handling non-ASCII characters
110 f.write('\n')