Update app.py

app.py

@@ -54,6 +54,12 @@ 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 > 4:
+            raise ValueError("max_new_tokens cannot be greater than 4.")
+        return v
+
 class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
         self.bucket_name = bucket_name
@@ -130,9 +136,7 @@ async def generate(request: GenerateRequest):
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay):
-    # Get the maximum model input length
     max_model_length = model.config.max_position_embeddings
-
     encoded_input = tokenizer(input_text, return_tensors="pt", max_length=max_model_length, truncation=True).to(device)
 
     def stop_criteria(input_ids, scores):
@@ -164,7 +168,7 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
             )
         except IndexError as e:
             print(f"IndexError during generation: {e}")
-            break  # Exit the loop if there's an index error
+            break
 
         new_token_ids = outputs.sequences[0][encoded_input.input_ids.shape[-1]:]