Update model.py

model.py

@@ -7,11 +7,11 @@ class ChestGPTDemo:
     def __init__(self, device=None):
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
 
-        # 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()
+        # Tiny GPT-2 model (lightweight and public)
+        self.tokenizer = AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2")
+        self.lm = AutoModelForCausalLM.from_pretrained("sshleifer/tiny-gpt2").to(self.device).eval()
 
-        # Few-shot prompt to guide generation
+        # Few-shot prompt to simulate clinical logic
         self.prompt = (
             "[radiology] Example 1:\n"
             "Global Disease: Cardiomegaly\n"
@@ -23,11 +23,11 @@ class ChestGPTDemo:
         )
 
     def process_image(self, img: Image.Image):
-        # Placeholder for image features – will add ViT later
+        # Placeholder for image encoding — will integrate ViT later
         return None
 
     def predict(self, img: Image.Image):
-        _ = self.process_image(img)  # skip vision for now
+        _ = self.process_image(img)
         inputs = self.tokenizer(self.prompt, return_tensors="pt", padding=True).to(self.device)
         outputs = self.lm.generate(
             **inputs,