@@ -33,6 +33,7 @@ def __init__(self,
3333 top_p : float = 1. ,
3434 temp : float = 1.0 ,
3535 repeat_penalty : float = 1 ,
36+ init_break : bool = True ,
3637 instruct_inp_prefix : str = "\n \n ### Instruction:\n \n " ,
3738 instruct_inp_suffix : str = "\n \n ### Response:\n \n " ,
3839 ) -> None :
@@ -48,6 +49,7 @@ def __init__(self,
4849 self .top_p = top_p
4950 self .temp = temp
5051 self .repeat_penalty = repeat_penalty
52+ self .init_break = init_break
5153
5254 # runtime args
5355 self .input_consumed = 0
@@ -81,9 +83,6 @@ def __init__(self,
8183 if (len (primer ) > 0 ):
8284 self .embd_inp += self ._tokenize (primer )
8385
84- # break immediately if using instruct
85- self .init_break = self .instruct
86-
8786 # number of tokens to keep when resetting context
8887 if (self .n_keep < 0 or self .n_keep > len (self .embd_inp ) or self .instruct ):
8988 self .n_keep = len (self .embd_inp )
@@ -182,13 +181,14 @@ def generate(self):
182181 if (len (self .embd_inp ) <= self .input_consumed ):
183182 # if antiprompt is present, stop
184183 if (self .use_antiprompt ()):
185- for i in self .first_antiprompt :
186- if i == self .last_n_tokens [- len (i ):]:
187- return
184+ if True in [
185+ i == self .last_n_tokens [- len (i ):]
186+ for i in self .first_antiprompt
187+ ]:
188+ break
188189
189190 # if we are using instruction mode, and we have processed the initial prompt
190191 if (self .init_break ):
191- self .init_break = False
192192 break
193193
194194 # if end of generation
@@ -201,6 +201,8 @@ def generate(self):
201201 self .embd_inp += self .first_antiprompt [0 ]
202202 break
203203
204+ self .init_break = False
205+
204206 def __enter__ (self ):
205207 return self
206208
0 commit comments