Update model.py

model.py

@@ -1,30 +1,37 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from torchvision import transforms
-from vit_encoder import ViTEncoder
 from PIL import Image
 
 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-rw-1b")
-        self.lm = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b").to(self.device).eval()
 
+        # Load clinical GPT-2 model
+        self.tokenizer = AutoTokenizer.from_pretrained("mrm8488/GPT-2-finetuned-clinical-notes")
+        self.lm = AutoModelForCausalLM.from_pretrained("mrm8488/GPT-2-finetuned-clinical-notes").to(self.device).eval()
+
+        # Few-shot prompt to guide generation
         self.prompt = (
-            "[radiology] Please describe this chest X-ray. "
-            "Mention global diseases and local findings if visible."
+            "[radiology] Example 1:\n"
+            "Global Disease: Cardiomegaly\n"
+            "Local Finding: Patchy opacity in right lower lobe (BBox: 50,60,120,150)\n\n"
+            "[radiology] Example 2:\n"
+            "Global Disease: Normal\n"
+            "Local Finding: No abnormalities detected\n\n"
+            "[radiology] Please describe this chest X-ray. Mention global diseases and local findings if visible.\n"
         )
 
     def process_image(self, img: Image.Image):
-        transform = transforms.Compose([
-            transforms.Resize((224, 224)),
-            transforms.ToTensor()
-        ])
-        return transform(img.convert("RGB")).unsqueeze(0).to(self.device)
+        # Placeholder for image features – will add ViT later
+        return None
 
     def predict(self, img: Image.Image):
-        _ = 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)
+        _ = self.process_image(img)  # skip vision for now
+        inputs = self.tokenizer(self.prompt, return_tensors="pt", padding=True).to(self.device)
+        outputs = self.lm.generate(
+            **inputs,
+            max_new_tokens=100,
+            pad_token_id=self.tokenizer.eos_token_id
+        )
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)