Create model.py

model.py

@@ -0,0 +1,38 @@
+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="cpu"):
+        self.device = device
+        self.vit = ViTEncoder().to(device).eval()
+        self.tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b-instruct")
+        self.lm = AutoModelForCausalLM.from_pretrained(
+            "tiiuae/falcon-7b-instruct",
+            device_map="auto",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        )
+        self.prompt = (
+            "[radiology] Please describe this chest X-ray in detail. "
+            "List global diseases and any local findings with locations."
+        )
+
+    def process_image(self, img: Image.Image):
+        transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor()
+        ])
+        tensor = transform(img.convert("RGB")).unsqueeze(0).to(self.device)
+        return tensor
+
+    def predict(self, img: Image.Image):
+        img_tensor = self.process_image(img)
+        with torch.no_grad():
+            vit_feat = self.vit(img_tensor)
+        # Shorten token list for demo
+        prompt = self.prompt + "\n[image_features]: " + ", ".join([f"{x:.3f}" for x in vit_feat[0][:10]])
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+        out = self.lm.generate(input_ids, max_new_tokens=100)
+        return self.tokenizer.decode(out[0], skip_special_tokens=True)