diff --git a/src/encoding.rs b/src/encoding.rs index 60257e7..febac72 100644 --- a/src/encoding.rs +++ b/src/encoding.rs @@ -1059,6 +1059,8 @@ pub struct StreamableParser { messages: Vec, state: StreamState, stop_tokens: HashSet, + start_token: Rank, + message_token: Rank, last_content_delta: Option, undecoded_tokens: Vec, undecoded_bytes: Vec, @@ -1090,6 +1092,8 @@ impl StreamableParser { options: ParseOptions, ) -> anyhow::Result { let stop_tokens = encoding.stop_tokens()?; + let start_token = encoding.render_formatting_token(FormattingToken::Start)?; + let message_token = encoding.render_formatting_token(FormattingToken::Message)?; let (state, next_role) = match role { Some(role) => ( StreamState::Header { @@ -1106,6 +1110,8 @@ impl StreamableParser { messages: Vec::new(), state, stop_tokens, + start_token, + message_token, last_content_delta: None, undecoded_tokens: Vec::new(), undecoded_bytes: Vec::new(), @@ -1123,9 +1129,7 @@ impl StreamableParser { let next_role_clone = self.next_role.clone(); match &mut self.state { StreamState::ExpectStart => { - let start = self - .encoding - .render_formatting_token(FormattingToken::Start)?; + let start = self.start_token; match token { Some(token) if token == start => { self.state = StreamState::Header { @@ -1147,9 +1151,7 @@ impl StreamableParser { } } StreamState::Header { header_tokens } => { - let msg_tok = self - .encoding - .render_formatting_token(FormattingToken::Message)?; + let msg_tok = self.message_token; match token { Some(token) if token == msg_tok => { // Clone the tokens and next_role, then clear the state before parsing