Upload model

CXR_LLAVA_HF.py

@@ -559,13 +559,42 @@ class CXRLLAVAModel(PreTrainedModel):
             raise ValueError(f'Unsupported tensor type: {return_tensors}')
         return input_ids
 
-    def generate_cxr_repsonse(self, chat, pil_image, temperature=0.2, top_p=0.8):
+    def write_radiologic_report(self, image, temperature=0.2, top_p=0.8):
+        chat = [
+            {"role": "system",
+             "content": "You are a helpful radiologist. Try to interpret chest x ray image and answer to the question that user provides."},
+            {"role": "user",
+             "content": "<image>\nWrite a radiologic report on the given chest radiograph, including information about atelectasis, cardiomegaly, consolidation, pulmonary edema, pleural effusion, and pneumothorax.\n"}
+        ]
+        response = self.generate_cxr_repsonse(chat=chat,image=image, temperature=temperature, top_p=top_p)
+        return response
+
+    def write_differential_diagnosis(self, image, temperature=0.2, top_p=0.8):
+        chat = [
+            {"role": "system",
+             "content": "You are a helpful radiologist. Try to interpret chest x ray image and answer to the question that user provides."},
+            {"role": "user",
+             "content": "<image>\nWhat are the possible differential diagnoses for this patient?\n"}
+        ]
+        response = self.generate_cxr_repsonse(chat=chat, image=image, temperature=temperature, top_p=top_p)
+        return response
+    def ask_question(self, question, image, temperature=0.2, top_p=0.8):
+        chat = [
+            {"role": "system",
+             "content": "You are a helpful radiologist. Try to interpret chest x ray image and answer to the question that user provides."},
+            {"role": "user",
+             "content": "<image>\n"+question}
+        ]
+        response = self.generate_cxr_repsonse(chat=chat, image=image, temperature=temperature, top_p=top_p)
+        return response
+
+    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=15)
             import numpy as np
-            pil_image = np.expand_dims(pil_image,axis=-1)
+            image = np.expand_dims(image,axis=-1)
             prompt = self.apply_chat_template(chat)
-            images = self.vision_tower.image_processor(pil_image, return_tensors='pt')['pixel_values']
+            images = self.vision_tower.image_processor(image, return_tensors='pt')['pixel_values']
             images = images.to(self.device)
             input_ids = self.tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
             stopping_criteria = KeywordsStoppingCriteria(["</s>"], self.tokenizer, input_ids)

VisualTransformer.py

@@ -1,3 +1,8 @@
+'''
+Source code from OPEN_CLIP project.
+https://github.com/mlfoundations/open_clip/blob/main/LICENSE
+'''
+
 from collections import OrderedDict
 import math
 from typing import Callable, Optional, Sequence, Tuple