| 15 | |
| 16 | |
| 17 | class APICallPostprocessing: |
| 18 | def __init__( |
| 19 | self, |
| 20 | start_tokens: List[int], |
| 21 | end_tokens: List[int], |
| 22 | minimum_percentage: float = 0.1, |
| 23 | ): |
| 24 | """ |
| 25 | Base API Postprocesing class |
| 26 | |
| 27 | :param start_tokens: token representation for [ or other tokens |
| 28 | :param end_tokens: token representation for ] or other tokens |
| 29 | :param minimum_percentage: pass percentage for candidate generation, less than this are ignored. |
| 30 | """ |
| 31 | self.start_tokens = start_tokens |
| 32 | self.end_tokens = end_tokens |
| 33 | self.minimum_percentage = minimum_percentage |
| 34 | self.api_text = "" # API text, might be better to pass it in |
| 35 | self.k_values = 5 # Default topk generation, might be better to pass it in |
| 36 | |
| 37 | def filter_continuations( |
| 38 | self, |
| 39 | input_tokens: torch.Tensor, |
| 40 | input_logits: torch.Tensor, |
| 41 | labels: torch.Tensor, |
| 42 | input_start: int, |
| 43 | tokenizer: PreTrainedTokenizerBase, |
| 44 | ) -> (torch.Tensor, torch.Tensor): |
| 45 | """ |
| 46 | Grab continuations that are valid |
| 47 | |
| 48 | :param input_tokens: tokenized inputs |
| 49 | :param input_logits: input logits |
| 50 | :param labels: labels for input logits |
| 51 | :param input_start: start of real input |
| 52 | :param tokenizer: |
| 53 | :return: Values and Indices |
| 54 | """ |
| 55 | # First, figure out locations... |
| 56 | probs = torch.softmax(input_logits, dim=-1) |
| 57 | # Make sure we don't keep any tokens that are supposed to be [ |
| 58 | remove_tokens = 1.0 - torch.sum( |
| 59 | torch.stack([labels == start_token for start_token in self.start_tokens]), |
| 60 | dim=0, |
| 61 | ) |
| 62 | # Get maximum probability... Should be sufficient. Maybe switch to sum if there's issues later |
| 63 | max_start_tokens = torch.amax( |
| 64 | torch.stack( |
| 65 | [probs[:, :, start_token] for start_token in self.start_tokens] |
| 66 | ), |
| 67 | dim=0, |
| 68 | ) |
| 69 | max_start_tokens = max_start_tokens * remove_tokens |
| 70 | return torch.topk(max_start_tokens[:, : -(M + 1)], k=self.k_values, dim=1) |
| 71 | |
| 72 | def create_candidates( |
| 73 | self, |
| 74 | indices: torch.Tensor, |
nothing calls this directly
no outgoing calls
no test coverage detected