ros/LLaMA-Factory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

tiny fix

Browse Source 
Former-commit-id: 81ef7017a4c96441951adeff0276cc5ab76a3544
This commit is contained in:
hiyouga
2023-08-03 17:42:28 +08:00
parent 27f4317ec6
commit d6e922dc1c
2 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/llmtuner/extras/misc.py
View File

@@ -29,7 +29,6 @@ class AverageMeter:
self.avg = self.sum / self.count
# Avoids runtime error in model.generate(do_sample=True).
class InvalidScoreLogitsProcessor(LogitsProcessor):
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
@@ -55,10 +54,10 @@ class StopWordsCriteria(StoppingCriteria):
return any([stop_id in input_ids[:, -1] for stop_id in self.stop_ids])
def get_stopwords_criteria(stop_ids: List[int]) -> StoppingCriteriaList:
stopwords_criteria = StoppingCriteriaList()
stopwords_criteria.append(StopWordsCriteria(stop_ids))
return stopwords_criteria
def get_stopping_criteria(stop_ids: List[int]) -> StoppingCriteriaList:
stopping_criteria = StoppingCriteriaList()
stopping_criteria.append(StopWordsCriteria(stop_ids))
return stopping_criteria
def count_parameters(model: torch.nn.Module) -> Tuple[int, int]: