(
request: Dict[str, Any],
update_stop_seq_fn: Callable[[List[str]], Tuple[List[List[int]], List[int]]],
)
| 27 | |
| 28 | |
| 29 | def process_stop_token_ids( |
| 30 | request: Dict[str, Any], |
| 31 | update_stop_seq_fn: Callable[[List[str]], Tuple[List[List[int]], List[int]]], |
| 32 | ) -> None: |
| 33 | stop_token_ids_final = [] |
| 34 | |
| 35 | if request.get("stop_token_ids") is not None: |
| 36 | stop_token_ids = request.get("stop_token_ids") |
| 37 | if isinstance(stop_token_ids, list) and len(stop_token_ids) > 0: |
| 38 | if isinstance(stop_token_ids[0], int): |
| 39 | # List[int] -> List[List[int]] |
| 40 | stop_token_ids_final.extend([[t] for t in stop_token_ids]) |
| 41 | elif isinstance(stop_token_ids[0], list): |
| 42 | # Already List[List[int]] |
| 43 | stop_token_ids_final.extend(stop_token_ids) |
| 44 | |
| 45 | stop_sequences = request.get("stop", []) |
| 46 | if stop_sequences: |
| 47 | stop_seqs, _ = update_stop_seq_fn(stop_sequences) |
| 48 | stop_token_ids_final.extend(stop_seqs) |
| 49 | |
| 50 | # Update request |
| 51 | if stop_token_ids_final: |
| 52 | stop_seqs_len = [len(seq) for seq in stop_token_ids_final] |
| 53 | request["stop_token_ids"] = stop_token_ids_final |
| 54 | request["stop_seqs_len"] = stop_seqs_len |
no test coverage detected