Update model.py

model.py

@@ -8,14 +8,12 @@ class ChestGPTDemo:
     def __init__(self, device=None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.vit = ViTEncoder().to(self.device).eval()
-        
-        self.tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
-        self.lm = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-7b-instruct")
-        self.lm = self.lm.to(self.device).eval()
+        self.tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b")
+        self.lm = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b").to(self.device).eval()
 
         self.prompt = (
-            "[radiology] Please describe this chest X-ray in detail. "
-            "List global diseases and any local findings with locations."
+            "[radiology] Please describe this chest X-ray. "
+            "Mention global diseases and local findings if visible."
         )
 
     def process_image(self, img: Image.Image):
@@ -23,18 +21,10 @@ class ChestGPTDemo:
             transforms.Resize((224, 224)),
             transforms.ToTensor()
         ])
-        tensor = transform(img.convert("RGB")).unsqueeze(0).to(self.device)
-        return tensor
+        return transform(img.convert("RGB")).unsqueeze(0).to(self.device)
 
     def predict(self, img: Image.Image):
-        img_tensor = self.process_image(img)
-        with torch.no_grad():
-            vit_feat = self.vit(img_tensor)
-
-        # Use first 10 features from ViT just as a mock
-        prompt = self.prompt + "\n[image_features]: " + ", ".join([f"{x:.3f}" for x in vit_feat[0][:10]])
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
-
-        output_ids = self.lm.generate(**inputs, max_new_tokens=100)
-        result = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
-        return result
+        _ = self.vit(self.process_image(img))  # placeholder for visual reasoning
+        input_ids = self.tokenizer(self.prompt, return_tensors="pt").input_ids.to(self.device)
+        output = self.lm.generate(input_ids, max_new_tokens=100)
+        return self.tokenizer.decode(output[0], skip_special_tokens=True)