Views
No views yet
1from transformers import (
2 AutoModelForCausalLM,
3 AutoTokenizer,
4 BitsAndBytesConfig,
5 TrainingArguments,
6 logging,
7)
8from peft import (
9 LoraConfig,
10 PeftModel,
11 get_peft_model,
12)
13import os, torch, gc
14from datasets import load_dataset
15import bitsandbytes as bnb
16from trl import SFTTrainer
17
18HF_TOKEN = "your-token"
19model_name = "llm-jp-3-13b-ry-ft"
20# QLoRA config
21
22bnb_config = BitsAndBytesConfig(
23 load_in_4bit=True,
24 bnb_4bit_quant_type="nf4",
25 bnb_4bit_compute_dtype=torch.bfloat16,
26)
27# Load model
28model = AutoModelForCausalLM.from_pretrained(
29 model_name,
30 quantization_config=bnb_config,
31 device_map="auto",
32 token = HF_TOKEN
33)
34# Load tokenizer
35tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, 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 = ""
46results = []
47for data in tqdm(datasets):
48 input = data["input"]
49
50 prompt = f"""### 指示
51 {input}
52 ### 回答：
53 """
54
55 tokenized_input = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
56 with torch.no_grad():
57 outputs = model.generate(
58 tokenized_input,
59 max_new_tokens=300,
60 do_sample=False,
61 repetition_penalty=1.2
62 )[0]
63 output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)
64 results.append({"task_id": data["task_id"], "input": input, "output": output})
65
66import re
67model_name = re.sub(".*/", "", model_name)
68with open(f"./{model_name}-my-original-outputs.jsonl", 'w', encoding='utf-8') as f:
69 for result in results:
70 json.dump(result, f, ensure_ascii=False) # ensure_ascii=False for handling non-ASCII characters
71 f.write('\n')