| 58 | |
| 59 | |
| 60 | class KeyWordsCriteria(StoppingCriteria): |
| 61 | def __init__(self, stop_id_sequences): |
| 62 | assert isinstance(stop_id_sequences[0], list), "stop_id_sequences should be a list of list of ids" |
| 63 | self.stop_sequences = stop_id_sequences |
| 64 | |
| 65 | def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: |
| 66 | sequences_should_be_stopped = [] |
| 67 | for i in range(input_ids.shape[0]): |
| 68 | sequence_should_be_stopped = False |
| 69 | for stop_sequence in self.stop_sequences: |
| 70 | if input_ids[i][-len(stop_sequence):].tolist() == stop_sequence: |
| 71 | sequence_should_be_stopped = True |
| 72 | break |
| 73 | sequences_should_be_stopped.append(sequence_should_be_stopped) |
| 74 | return all(sequences_should_be_stopped) |
| 75 | |
| 76 | |
| 77 | @torch.no_grad() |
nothing calls this directly
no outgoing calls
no test coverage detected