Update app.py

app.py

@@ -70,6 +70,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 > 10:
+            raise ValueError("max_new_tokens cannot exceed 10.")
+        return v
+
 class GCSModelLoader:
     def __init__(self, bucket):
         self.bucket = bucket
@@ -119,7 +125,7 @@ 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
+    initial_max_new_tokens = request.max_new_tokens  # The requested max tokens (will be max 10)
     generation_params = request.model_dump(
         exclude_none=True,
         exclude={'model_name', 'input_text', 'task_type', 'stream', 'chunk_delay', 'max_new_tokens'}
@@ -133,18 +139,38 @@ async def generate(request: GenerateRequest):
         tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
 
         async def generate_responses() -> AsyncIterator[Dict[str, List[Dict[str, str]]]]:
-            text_pipeline = pipeline(task_type, model=model_name, tokenizer=tokenizer, token=HUGGINGFACE_HUB_TOKEN, **generation_params, max_new_tokens=initial_max_new_tokens)
-
-            def generate_on_thread(pipeline, input_text, output_queue):
-                result = pipeline(input_text)
-                output_queue.put_nowait(result)
-
-            output_queue = asyncio.Queue()
-            thread = Thread(target=generate_on_thread, args=(text_pipeline, input_text, output_queue))
-            thread.start()
-            result = await output_queue.get()
-            thread.join()
-            yield {"response": result}
+            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)
+                text_pipeline = pipeline(
+                    task_type,
+                    model=model_name,
+                    tokenizer=tokenizer,
+                    token=HUGGINGFACE_HUB_TOKEN,
+                    **generation_params,
+                    max_new_tokens=current_max_new_tokens
+                )
+
+                def generate_on_thread(pipeline, input_text, output_queue):
+                    result = pipeline(input_text)
+                    output_queue.put_nowait(result)
+
+                output_queue = asyncio.Queue()
+                thread = Thread(target=generate_on_thread, args=(text_pipeline, input_text, output_queue))
+                thread.start()
+                result = await output_queue.get()
+                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
+
+                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
+                input_text = all_generated_text
 
         async def text_stream():
             async for data in generate_responses():