Update README.md

README.md

@@ -124,9 +124,8 @@ from transformers import (
 from peft import LoraConfig, PeftModel
 from datasets import load_dataset
 
-
 BASE_MODEL = "llm-jp/llm-jp-3-13b"
-PEFT_MODEL = "libkazz/llm-jp-3-13b-it"
+PEFT_MODEL = "libkazz/llm-jp-3-13b-it_lora"
 
 tokenizer = AutoTokenizer.from_pretrained(PEFT_MODEL)
 bnb_config = BitsAndBytesConfig(
@@ -148,30 +147,32 @@ model = PeftModel.from_pretrained(base_model, PEFT_MODEL)
 
 # elyza-tasks-100-TV_0.jsonl データの読み込み
 from datasets import load_dataset
+from tqdm import tqdm
 
-dataset = load_dataset("json", data_files="./elyza-tasks-100-TV_0.jsonl", split="train")
+datasets = load_dataset("json", data_files="./elyza-tasks-100-TV_0.jsonl", split="train")
 
 results = []
-for num in tqdm(range(100)):
-    instruction = dataset["input"][num]
-
-    prompt = f"次の指示に忠実に回答を作成しなさい。\n\n### 指示:\n{instruction}\n\n### 回答:\n"
-
-    model_input = tokenizer(prompt, return_tensors="pt").to(model.device)
-    input_ids = model_input["input_ids"]
-
-    with torch.no_grad():
-        outputs = model.generate(
-            input_ids,
-            max_new_tokens=300,
-            attention_mask = model_input.attention_mask,
-            pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            do_sample=False,
-            repetition_penalty=1.02,
-        )[0]
-    output = tokenizer.decode(outputs[input_ids.size(1):], skip_special_tokens=True)
-    results.append({"task_id": num, "input": instruction, "output": output})
+for dt in tqdm(datasets):
+  input = dt["input"]
+  prompt = f"次の指示に忠実に回答を作成しなさい。\n\n### 指示:\n{input}\n\n### 回答:\n"
+
+  model_input = tokenizer([prompt], return_tensors = "pt").to(model.device)
+
+  if "token_type_ids" in model_input:
+      del model_input["token_type_ids"]
+  
+  outputs = model.generate(
+      **model_input,
+      max_new_tokens = 512,
+      pad_token_id=tokenizer.pad_token_id,
+      eos_token_id=tokenizer.eos_token_id,
+      use_cache = True,
+      do_sample=False,
+      repetition_penalty=1.2
+    )
+  
+  prediction = tokenizer.decode(outputs[0], skip_special_tokens=True).split('\n### 回答')[-1]
+  results.append({"task_id": dt["task_id"], "input": input, "output": prediction})
 
 # ファイルに保存する
 import json