diff --git a/modeling/inference_models/hf_torch.py b/modeling/inference_models/hf_torch.py index 1b411c95..50aa0538 100644 --- a/modeling/inference_models/hf_torch.py +++ b/modeling/inference_models/hf_torch.py @@ -225,9 +225,6 @@ class HFTorchInferenceModel(HFInferenceModel): ) class KoboldLogitsWarperList(LogitsProcessorList): - def __init__(self): - pass - def __call__( lw_self, input_ids: torch.LongTensor, @@ -244,16 +241,10 @@ class HFTorchInferenceModel(HFInferenceModel): ), f"Scores are None; processor '{processor}' is to blame" return scores - def new_get_logits_warper( - beams: int = 1, - ) -> LogitsProcessorList: - return KoboldLogitsWarperList() - def new_sample(self, *args, **kwargs): assert kwargs.pop("logits_warper", None) is not None - kwargs["logits_warper"] = new_get_logits_warper( - beams=1, - ) + kwargs["logits_warper"] = KoboldLogitsWarperList() + if utils.koboldai_vars.newlinemode in ["s", "ns"]: kwargs["eos_token_id"] = -1 kwargs.setdefault("pad_token_id", 2) diff --git a/modeling/logits_processors.py b/modeling/logits_processors.py index 20a18026..d3133109 100644 --- a/modeling/logits_processors.py +++ b/modeling/logits_processors.py @@ -267,6 +267,14 @@ class PhraseBiasLogitsProcessor: for batch in range(scores_shape[0]): for token, bias in self._get_biased_tokens(input_ids[batch]).items(): - scores[batch][token] += bias + if bias > 0 and bool(scores[batch][token].isneginf()): + # Adding bias to -inf will do NOTHING!!! So just set it for + # now. There may be more mathishly correct way to do this + # but it'll work. Also, make sure the bias is actually + # positive. Don't give a -inf token more chance by setting + # it to -0.5! + scores[batch][token] = bias + else: + scores[batch][token] += bias return scores