Update app.py

app.py

@@ -74,9 +74,10 @@ class GenerateRequest(BaseModel):
         return v
 
 class StopOnKeywords(StoppingCriteria):
-    def __init__(self, stop_words_ids: List[List[int]], encounters: int = 1):
+    def __init__(self, stop_words_ids: List[List[int]], tokenizer, encounters: int = 1):
         super().__init__()
         self.stop_words_ids = stop_words_ids
+        self.tokenizer = tokenizer
         self.encounters = encounters
         self.current_encounters = 0
 
@@ -135,7 +136,7 @@ model_loader = GCSModelLoader(bucket)
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     model_name = request.model_name
-    input_text = request.input_text  # Initialize input_text here
+    input_text = request.input_text
     task_type = request.task_type
     requested_max_new_tokens = request.max_new_tokens
     generation_params = request.model_dump(
@@ -155,7 +156,7 @@ async def generate(request: GenerateRequest):
 
         if user_defined_stopping_strings:
             stop_words_ids = [tokenizer.encode(stop_string, add_special_tokens=False) for stop_string in user_defined_stopping_strings]
-            stopping_criteria_list.append(StopOnKeywords(stop_words_ids))
+            stopping_criteria_list.append(StopOnKeywords(stop_words_ids, tokenizer)) # Pass tokenizer
 
         if config.eos_token_id is not None:
             eos_token_ids = [config.eos_token_id]
@@ -164,13 +165,13 @@ async def generate(request: GenerateRequest):
             elif isinstance(config.eos_token_id, list):
                 eos_token_ids = [[id] for id in config.eos_token_id]
             stop_words_ids_eos = [tokenizer.encode(tokenizer.decode(eos_id), add_special_tokens=False) for eos_id in eos_token_ids]
-            stopping_criteria_list.append(StopOnKeywords(stop_words_ids_eos))
+            stopping_criteria_list.append(StopOnKeywords(stop_words_ids_eos, tokenizer)) # Pass tokenizer
         elif tokenizer.eos_token is not None:
             stop_words_ids_eos = [tokenizer.encode(tokenizer.eos_token, add_special_tokens=False)]
-            stopping_criteria_list.append(StopOnKeywords(stop_words_ids_eos))
+            stopping_criteria_list.append(StopOnKeywords(stop_words_ids_eos, tokenizer)) # Pass tokenizer
 
         async def generate_responses() -> AsyncIterator[Dict[str, List[Dict[str, str]]]]:
-            nonlocal input_text  # Allow modification of the outer scope variable
+            nonlocal input_text
             all_generated_text = ""
             stop_reason = None
 
@@ -195,22 +196,25 @@ async def generate(request: GenerateRequest):
                 result = await output_queue.get()
                 thread.join()
 
-                newly_generated_text = result[0]['generated_text'][len(all_generated_text):]
-
-                if not newly_generated_text:
+                newly_generated_text = result[0]['generated_text']
+                
+                # Decode tokens to check for stopping strings
+                for criteria in stopping_criteria_list:
+                    if isinstance(criteria, StopOnKeywords):
+                        for stop_ids in criteria.stop_words_ids:
+                            decoded_stop_string = tokenizer.decode(stop_ids)
+                            if decoded_stop_string in newly_generated_text:
+                                stop_reason = f"stopping_string: {decoded_stop_string}"
+                                break
+                        if stop_reason:
+                            break
+                
+                if stop_reason:
                     break
 
                 all_generated_text += newly_generated_text
                 yield {"response": [{'generated_text': newly_generated_text}]}
 
-                if stopping_criteria_list:
-                    for criteria in stopping_criteria_list:
-                        if isinstance(criteria, StopOnKeywords) and criteria.current_encounters > 0:
-                            stop_reason = "stopping_string"
-                            break
-                    if stop_reason:
-                        break
-
                 if config.eos_token_id is not None:
                     eos_tokens = [config.eos_token_id]
                     if isinstance(config.eos_token_id, int):