423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-16 03:40:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

implement stream generating

Browse Source
This commit is contained in:
hiyouga
2023-06-05 16:43:44 +08:00
parent 44298c1235
commit fe1d930816
3 changed files with 32 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/utils/other.py
View File

@@ -52,13 +52,12 @@ class AverageMeter:
# Avoid runtime error in model.generate(do_sample=True).
# Borrowed from: https://huggingface.co/THUDM/chatglm-6b/blob/658202d88ac4bb782b99e99ac3adff58b4d0b813/modeling_chatglm.py#L54
class InvalidScoreLogitsProcessor(LogitsProcessor):
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
if torch.isnan(scores).any() or torch.isinf(scores).any():
scores.zero_()
scores[..., 5] = 5e4
scores[:, 0] = 1.0
return scores