Update app.py

app.py

@@ -191,11 +191,13 @@ async def generate_stream(model, tokenizer, input_text,
     async def stream():
          past_key_values = None
          input_ids = None
-         async for token,past_key_values_response,input_ids_response in stream_text(model, tokenizer, input_text,
+         async for token,past_key_values_response,input_ids_response, is_end in stream_text(model, tokenizer, input_text,
                              generation_config, stop_sequences,
                              device,pad_token_id, max_model_length, max_new_tokens, past_key_values, input_ids):
             past_key_values = past_key_values_response
             input_ids = input_ids_response
+            if is_end:
+                break
             yield token
     return stream()
 
@@ -283,25 +285,24 @@ async def stream_text(model, tokenizer, input_text,
     
     output_text = ""
     stop_criteria = StoppingCriteriaList([StopOnSequencesCriteria(stop_sequences, tokenizer)])
-
-    while True:
-        
-        if input_ids is None:
-             encoded_input = tokenizer(
+    
+    if input_ids is None:
+      encoded_input = tokenizer(
                 input_text, return_tensors="pt",
                 truncation=True,
                 padding = "max_length",
                  max_length=max_model_length
             ).to(device)
-             input_ids = encoded_input.input_ids
-        else:
-          encoded_input = {
-          "input_ids":input_ids,
-          "past_key_values": past_key_values
-        }
-        
+      input_ids = encoded_input.input_ids
+    else:
+      encoded_input = {
+       "input_ids":input_ids,
+       "past_key_values": past_key_values
+      }
+    
+    while True:
         outputs = model.generate(
-             **encoded_input,
+            **encoded_input,
             do_sample=generation_config.do_sample,
             max_new_tokens=generation_config.max_new_tokens,
             temperature=generation_config.temperature,
@@ -313,7 +314,7 @@ async def stream_text(model, tokenizer, input_text,
             return_dict_in_generate=True,
             pad_token_id=pad_token_id if pad_token_id is not None else None,
             stopping_criteria = stop_criteria,
-           
+            
         )
         
         new_text = tokenizer.decode(
@@ -324,17 +325,19 @@ async def stream_text(model, tokenizer, input_text,
         output_text += new_text
 
         stop_index = find_stop(output_text, stop_sequences)
-
-        if stop_index != -1:
-            final_output = output_text[:stop_index]
+        
+        is_end = False
+        if stop_index != -1 or (hasattr(outputs, "sequences") and outputs.sequences[0][-1] == tokenizer.eos_token_id):
+            final_output = output_text[:stop_index] if stop_index != -1 else output_text
             
             for text in final_output.split():
                 yield json.dumps({"text": text, "is_end": False}) + "\n", \
                       outputs.past_key_values if hasattr(outputs, "past_key_values") else None , \
-                      outputs.sequences if hasattr(outputs, "sequences") else None
-            yield json.dumps({"text": "", "is_end": True}) + "\n", \
+                      outputs.sequences if hasattr(outputs, "sequences") else None, True
+            
+            yield json.dumps({"text": "", "is_end": True}) + "\n",\
                     outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
-                     outputs.sequences if hasattr(outputs, "sequences") else None
+                     outputs.sequences if hasattr(outputs, "sequences") else None, True
             break
         else:
            
@@ -345,19 +348,19 @@ async def stream_text(model, tokenizer, input_text,
                 chunk = tokens[i:i + max_new_tokens]
                 chunk_text = " ".join(chunk)
                 for text in chunk_text.split():
-                    yield json.dumps({"text": text, "is_end": False}) + "\n", \
+                     yield json.dumps({"text": text, "is_end": False}) + "\n", \
                         outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
-                         outputs.sequences if hasattr(outputs, "sequences") else None
+                         outputs.sequences if hasattr(outputs, "sequences") else None, False
         
         if len(new_text) == 0:
             
             for text in output_text.split():
                 yield json.dumps({"text": text, "is_end": False}) + "\n", \
                         outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
-                         outputs.sequences if hasattr(outputs, "sequences") else None
+                         outputs.sequences if hasattr(outputs, "sequences") else None, True
             yield json.dumps({"text": "", "is_end": True}) + "\n",\
                     outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
-                     outputs.sequences if hasattr(outputs, "sequences") else None
+                     outputs.sequences if hasattr(outputs, "sequences") else None, True
             break
             
         past_key_values = outputs.past_key_values if hasattr(outputs, "past_key_values") else None
@@ -366,7 +369,6 @@ async def stream_text(model, tokenizer, input_text,
         output_text = ""
 
 
-
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try: