Update app.py

app.py

@@ -130,32 +130,36 @@ 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):
-    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
+    encoded_input = tokenizer(input_text, return_tensors="pt").to(device)
 
     def stop_criteria(input_ids, scores):
-        decoded_output = tokenizer.decode(int(input_ids[0][-1]), skip_special_tokens=True)
-        return decoded_output in stop_sequences
+        decoded_output = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+        for stop in stop_sequences:
+            if decoded_output.endswith(stop):
+                return True
+        return False
 
     stopping_criteria = StoppingCriteriaList([stop_criteria])
 
     token_buffer = []
-    outputs = model.generate(
-        **encoded_input,
-        do_sample=generation_config.do_sample,
-        max_new_tokens=generation_config.max_new_tokens,
-        temperature=generation_config.temperature,
-        top_p=generation_config.top_p,
-        top_k=generation_config.top_k,
-        repetition_penalty=generation_config.repetition_penalty,
-        num_return_sequences=generation_config.num_return_sequences,
-        stopping_criteria=stopping_criteria,
-        output_scores=True,
-        return_dict_in_generate=True,
-        streamer=None  # Ensure streamer is None for manual token processing
-    )
-
-    for output in outputs.sequences:
-        for token_id in output:
+    output_ids = encoded_input.input_ids
+    while True:
+        outputs = model.generate(
+            output_ids,
+            do_sample=generation_config.do_sample,
+            max_new_tokens=generation_config.max_new_tokens,
+            temperature=generation_config.temperature,
+            top_p=generation_config.top_p,
+            top_k=generation_config.top_k,
+            repetition_penalty=generation_config.repetition_penalty,
+            num_return_sequences=generation_config.num_return_sequences,
+            stopping_criteria=stopping_criteria,
+            output_scores=True,
+            return_dict_in_generate=True
+        )
+        new_token_ids = outputs.sequences[0][encoded_input.input_ids.shape[-1]:]
+
+        for token_id in new_token_ids:
             token = tokenizer.decode(token_id, skip_special_tokens=True)
             token_buffer.append(token)
             if len(token_buffer) >= 10:
@@ -167,26 +171,13 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
             yield "".join(token_buffer)
             token_buffer = []
 
-        if stop_sequences and any(stop in tokenizer.decode(output, skip_special_tokens=True) for stop in stop_sequences):
-            return
+        if stop_criteria(outputs.sequences, None):
+            break
 
-        encoded_input = tokenizer.build_inputs_with_special_tokens(output)
-        encoded_input = {'input_ids': torch.tensor([encoded_input]).to(device)}
+        if len(new_token_ids) < generation_config.max_new_tokens:
+            break
 
-        outputs = model.generate(
-            **encoded_input,
-            do_sample=generation_config.do_sample,
-            max_new_tokens=generation_config.max_new_tokens,
-            temperature=generation_config.temperature,
-            top_p=generation_config.top_p,
-            top_k=generation_config.top_k,
-            repetition_penalty=generation_config.repetition_penalty,
-            num_return_sequences=generation_config.num_return_sequences,
-            stopping_criteria=stopping_criteria,
-            output_scores=True,
-            return_dict_in_generate=True,
-            streamer=None
-        )
+        output_ids = outputs.sequences
 
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):