[Enhancement] Extending the metaphor to regularization

[chrico-bu-uab]

This function pushes "weights" (prompts) to be "small" (short).

class TokenCountLoss(Module):
    def __init__(self, max_tokens):
        super().__init__()
        self.call_fn = StringBasedFunction(
            fn=partial(self._eval_df, max_tokens=max_tokens),
            function_purpose="evaluating the token count of the model's response",
        )

    @classmethod
    def _eval_df(cls, input_val: tg.Variable, max_tokens: int) -> str:
        tc = get_token_count(input_val.value)
        out = f"token count: {tc}; "
        if tc > max_tokens:
            return out + f"exceeded max tokens by {tc - max_tokens}"
        return out + "within limits"

    def forward(self, **kwargs):
        return self.call_fn(inputs=kwargs)

...
reg_loss = TokenCountLoss(max_tokens)(input_val=model.system_prompt)
losses.append(reg_loss)
total_loss = tg.sum(losses)