def test_beamsearch(n, m, k, s):
torch.manual_seed(s)
prob = torch.rand((n, m))
output = beamsearch(prob, k)
if n != 0:
assert output.shape == (min(m ** n, k), n)
assert all(output[0] == prob.argmax(1))
assert all(output[0] == beamsearch(prob, 1)[0])
def get_best_tags(logit: torch.Tensor, id2label: List[str], k_beam: int) -> List[str]:
"""Select best tags from logit based on beamsearch."""
candidates = beamsearch(logit.softmax(-1), k_beam)
assert len(cast(Sized, candidates))
best_tags: List[str] = []
for cand in candidates:
tags, is_correct = correct_bio_tags([id2label[j] for j in cast(Iterable, cand)])
if is_correct:
best_tags = tags
break
if not best_tags:
best_tags = tags
best_tags = [t if t != "-" else "O" for t in best_tags]
return best_tags