Update README.md

README.md

@@ -53,6 +53,52 @@ def load_model(model_name):
   return model, tokenizer
 
 
+def inference(datasets, model, tokenizer):
+  _results = []
+  for data in tqdm(datasets):
+      input = data["input"]
+
+      prompt = f"""### 指示
+      {input}
+      ### 回答：
+      """
+
+      # 修正箇所: encode_plus を使用して attention_mask を取得
+      encoded_input = tokenizer.encode_plus(
+          prompt,
+          add_special_tokens=False,
+          return_tensors="pt",
+          padding=True,
+          truncation=True,
+      ).to(model.device)
+
+      tokenized_input = encoded_input["input_ids"]
+      attention_mask = encoded_input["attention_mask"]
+
+      # 修正箇所: attention_mask と pad_token_id を model.generate に渡す
+      with torch.no_grad():
+          outputs = model.generate(
+              tokenized_input,
+              attention_mask=attention_mask,
+              max_new_tokens=100,
+              do_sample=False,
+              repetition_penalty=1.2,
+              pad_token_id=tokenizer.pad_token_id  # 明示的に設定
+          )[0]
+
+      output = tokenizer.decode(
+          outputs[tokenized_input.size(1):],
+          skip_special_tokens=True
+          )
+
+      _results.append({
+          "task_id": data["task_id"],
+          "input": input,
+          "output": output
+      })
+  return _results
+
+
 model_name = "ak0327/llm-jp-3-13b-ft-5"
 
 model, tokenizer = load_model(model_name)