From 0961f5a49a7f826a8b8be08451b308b374868c51 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 13 Feb 2023 02:39:12 +0000 Subject: [PATCH] Add find method to sequence group --- cacheflow/sequence.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/cacheflow/sequence.py b/cacheflow/sequence.py index 619ccd8b..69e84641 100644 --- a/cacheflow/sequence.py +++ b/cacheflow/sequence.py @@ -76,3 +76,9 @@ class SequenceGroup: return len(self.seqs) else: return len([seq for seq in self.seqs if seq.status == status]) + + def find(self, seq_id: int) -> Sequence: + for seq in self.seqs: + if seq.seq_id == seq_id: + return seq + raise ValueError(f'Sequence {seq_id} not found.')