github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 11:02:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[eval] update llama npu eval (#5366)

Browse Source
This commit is contained in:
Camille Zhong
2024-02-06 10:53:03 +08:00
committed by GitHub
parent 44ca61a22b
commit a5756a8720
3 changed files with 18 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
applications/ColossalEval/colossal_eval/models/chatglm.py
View File

@@ -3,6 +3,8 @@ from typing import List
import torch
from colossalai.utils import get_current_device
from .huggingface import HuggingFaceModel
IGNORE_INDEX = -100
@@ -126,9 +128,9 @@ class ChatGLMModel(HuggingFaceModel):
"""
input_ids = torch.nn.utils.rnn.pad_sequence(
input_ids_list, batch_first=True, padding_value=self.tokenizer.pad_token_id
).to(torch.cuda.current_device())
).to(get_current_device())
labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX).to(
torch.cuda.current_device()
get_current_device()
)
outputs = self.model(input_ids)[0]
@@ -197,7 +199,7 @@ class ChatGLM2Model(ChatGLMModel):
truncation=True,
return_tensors="pt",
max_length=self.model_max_length - max_new_tokens,
).to(torch.cuda.current_device())
).to(get_current_device())
# Set output_scores=True to get prediction scores.
outputs = self.model.generate(