Update app.py

app.py

@@ -14,6 +14,8 @@ from dotenv import load_dotenv
 import huggingface_hub
 from threading import Thread
 from typing import AsyncIterator, List, Dict
+from transformers.stopping_criteria import StoppingCriteria, StoppingCriteriaList
+import torch
 
 load_dotenv()
 
@@ -55,7 +57,8 @@ class GenerateRequest(BaseModel):
     num_return_sequences: int = 1
     do_sample: bool = False
     chunk_delay: float = 0.0
-    max_new_tokens: int = 512
+    max_new_tokens: int = 10
+    stopping_strings: List[str] = None
 
     @field_validator("model_name")
     def model_name_cannot_be_empty(cls, v):
@@ -70,11 +73,20 @@ class GenerateRequest(BaseModel):
             raise ValueError(f"task_type must be one of: {valid_types}")
         return v
 
-    @field_validator("max_new_tokens")
-    def max_new_tokens_must_be_within_limit(cls, v):
-        if v > 10:
-            raise ValueError("max_new_tokens cannot exceed 10.")
-        return v
+class StopOnKeywords(StoppingCriteria):
+    def __init__(self, stop_words_ids: List[List[int]], encounters: int = 1):
+        super().__init__()
+        self.stop_words_ids = stop_words_ids
+        self.encounters = encounters
+        self.current_encounters = 0
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        for stop_ids in self.stop_words_ids:
+            if torch.all(input_ids[0][-len(stop_ids):] == torch.tensor(stop_ids).to(input_ids.device)):
+                self.current_encounters += 1
+                if self.current_encounters >= self.encounters:
+                    return True
+        return False
 
 class GCSModelLoader:
     def __init__(self, bucket):
@@ -125,11 +137,12 @@ async def generate(request: GenerateRequest):
     model_name = request.model_name
     input_text = request.input_text
     task_type = request.task_type
-    initial_max_new_tokens = request.max_new_tokens  # The requested max tokens (will be max 10)
+    requested_max_new_tokens = request.max_new_tokens
     generation_params = request.model_dump(
         exclude_none=True,
-        exclude={'model_name', 'input_text', 'task_type', 'stream', 'chunk_delay', 'max_new_tokens'}
+        exclude={'model_name', 'input_text', 'task_type', 'stream', 'chunk_delay', 'max_new_tokens', 'stopping_strings'}
     )
+    user_defined_stopping_strings = request.stopping_strings
 
     try:
         if not model_loader.check_model_exists_locally(model_name):
@@ -137,19 +150,40 @@ async def generate(request: GenerateRequest):
                 raise HTTPException(status_code=500, detail=f"Failed to load model: {model_name}")
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
+        config = AutoConfig.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
+        stopping_criteria_list = StoppingCriteriaList()
+
+        # Add user-defined stopping strings if provided
+        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))
+
+        # Automatically add EOS token as a stopping criterion
+        if config.eos_token_id is not None:
+            eos_token_ids = [config.eos_token_id]
+            if isinstance(config.eos_token_id, int):
+                eos_token_ids = [[config.eos_token_id]]
+            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))
+        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))
 
         async def generate_responses() -> AsyncIterator[Dict[str, List[Dict[str, str]]]]:
             all_generated_text = ""
-            remaining_tokens = 512 # Or some reasonable maximum
-            while remaining_tokens > 0:
-                current_max_new_tokens = min(initial_max_new_tokens, remaining_tokens)
+            stop_reason = None  # To track why the generation stopped
+
+            while True: # Loop indefinitely, relying on stopping criteria
                 text_pipeline = pipeline(
                     task_type,
                     model=model_name,
                     tokenizer=tokenizer,
                     token=HUGGINGFACE_HUB_TOKEN,
+                    stopping_criteria=stopping_criteria_list,
                     **generation_params,
-                    max_new_tokens=current_max_new_tokens
+                    max_new_tokens=requested_max_new_tokens  # Generate in chunks
                 )
 
                 def generate_on_thread(pipeline, input_text, output_queue):
@@ -163,13 +197,40 @@ async def generate(request: GenerateRequest):
                 thread.join()
 
                 newly_generated_text = result[0]['generated_text'][len(all_generated_text):]
-                if not newly_generated_text: # Break if no new text is generated
-                    break
+
+                if not newly_generated_text:
+                    break # Should ideally not happen with proper stopping criteria
 
                 all_generated_text += newly_generated_text
                 yield {"response": [{'generated_text': newly_generated_text}]}
-                remaining_tokens -= current_max_new_tokens
-                # Update input_text for the next iteration
+
+                # Check if any stopping criteria was met
+                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 the generated text seems to match the EOS token, stop
+                if config.eos_token_id is not None:
+                    eos_tokens = [config.eos_token_id]
+                    if isinstance(config.eos_token_id, int):
+                        eos_tokens = [config.eos_token_id]
+                    elif isinstance(config.eos_token_id, list):
+                        eos_tokens = config.eos_token_id
+                    for eos_token in eos_tokens:
+                        if tokenizer.decode([eos_token]) in newly_generated_text:
+                            stop_reason = "eos_token"
+                            break
+                    if stop_reason:
+                        break
+                elif tokenizer.eos_token is not None and tokenizer.eos_token in newly_generated_text:
+                    stop_reason = "eos_token"
+                    break
+
+                # Update input text for the next iteration
                 input_text = all_generated_text
 
         async def text_stream():
@@ -186,4 +247,5 @@ async def generate(request: GenerateRequest):
         raise HTTPException(status_code=500, detail=f"Internal server error: {e}")
 
 if __name__ == "__main__":
+    import torch
     uvicorn.run(app, host="0.0.0.0", port=7860)