add EOS state cache point
This commit is contained in:
		
							parent
							
								
									4339ce20d5
								
							
						
					
					
						commit
						16b636ef83
					
				@ -272,6 +272,17 @@ class AbstractRWKV(ABC):
 | 
			
		||||
            )
 | 
			
		||||
 | 
			
		||||
            if token == self.EOS_ID:
 | 
			
		||||
                try:
 | 
			
		||||
                    state_cache.add_state(
 | 
			
		||||
                        state_cache.AddStateBody(
 | 
			
		||||
                            prompt=prompt + response,
 | 
			
		||||
                            tokens=self.model_tokens,
 | 
			
		||||
                            state=self.model_state,
 | 
			
		||||
                            logits=logits,
 | 
			
		||||
                        )
 | 
			
		||||
                    )
 | 
			
		||||
                except HTTPException:
 | 
			
		||||
                    pass
 | 
			
		||||
                yield response, "", prompt_token_len, completion_token_len
 | 
			
		||||
                break
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user