Update CXR_LLAVA_HF.py

CXR_LLAVA_HF.py

@@ -10,6 +10,7 @@ from threading import Thread
 from dataclasses import dataclass
 import numpy as np
 from PIL import Image
+
 # Model Constants
 IGNORE_INDEX = -100
 IMAGE_TOKEN_INDEX = -200
@@ -597,7 +598,7 @@ class CXRLLAVAModel(PreTrainedModel):
 
     def generate_cxr_repsonse(self, chat, image, temperature=0.2, top_p=0.8):
         with torch.no_grad():
-            streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=90)
+            streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=180)
 
             if np.array(image).max()>255:
                 raise Exception("16-bit image is not supported.")
@@ -636,8 +637,11 @@ class CXRLLAVAModel(PreTrainedModel):
             ))
             thread.start()
             generated_text = ""
+            text_len = 0
             for new_text in streamer:
                 generated_text += new_text
+                text_len += 1
+                if text_len > 200: break
 
         return generated_text