# Source: content/notes/nlp/text-generation-and-decoding.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import json
import numpy as np
from transformers import RepetitionPenaltyLogitsProcessor
import torch
torch.set_num_threads(1)
logits = torch.tensor([[-2., -1., 0., 1., 2.]])
seen = torch.tensor([[0, 4]])
penalized = RepetitionPenaltyLogitsProcessor(1.2)(seen, logits.clone())
assert penalized[0, 0] < logits[0, 0] and penalized[0, 4] < logits[0, 4]
p = torch.softmax(penalized[0]/.8, -1).numpy()
order = np.argsort(-p, kind="stable")
k = int(np.searchsorted(np.cumsum(p[order]), .8, side="left"))+1
kept = order[:k]
assert p[kept].sum() >= .8 and (k == 1 or p[order[:k-1]].sum() < .8)
min_p = p >= .1*p.max()
assert min_p[p.argmax()]
try:
    json.loads('{"answer": "unfinished')
except json.JSONDecodeError:
    pass
else:
    raise AssertionError("incomplete JSON unexpectedly accepted")
print("penalized logits", penalized.tolist(), "nucleus", kept.tolist(), "min-p mask", min_p.tolist())