Upload app.py

app.py

@@ -262,10 +262,11 @@ def stream_chat(input_images: List[Image.Image], caption_type: str, caption_tone
         generate_ids = text_model.generate(input_ids=input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, do_sample=True, 
                                            suppress_tokens=None, max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature)
 
+        # Trim off the prompt
         generate_ids = generate_ids[:, input_ids.shape[1]:]
-
         for ids in generate_ids:
-            caption = tokenizer.decode(ids[:-1] if ids[-1] == tokenizer.eos_token_id else ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+            caption = tokenizer.decode(ids[:] if ids[0] == tokenizer.eos_token_id or ids[0] == tokenizer.convert_tokens_to_ids("<|eot_id|>") else ids,
+                                       skip_special_tokens=True, clean_up_tokenization_spaces=True)
             caption = caption.replace('<|end_of_text|>', '').replace('<|finetune_right_pad_id|>', '').strip()
             all_captions.append(caption)