Update app.py

app.py

@@ -1,26 +1,116 @@
-# app.p
+# app.py
 import gradio as gr
-from transformers import AutoTokenizer, AutoModel, AutoFeatureExtractor
 import torch
 from PIL import Image
+import json
+import os
+from tokenizers import ByteLevelBPETokenizer # Changed from Tokenizer
+from torchvision import transforms
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("hackergeek/radiology-image-captioning")
-model = AutoModel.from_pretrained("hackergeek/radiology-image-captioning")
-feature_extractor = AutoFeatureExtractor.from_pretrained("hackergeek/radiology-image-captioning")
+# ============================================================
+# 3. Model (Tiny CNN + Transformer Decoder) - Re-define model classes
+# ============================================================
+class CNNEncoder(torch.nn.Module):
+    def __init__(self, embed_dim=128):
+        super().__init__()
+        self.cnn = torch.nn.Sequential(
+            torch.nn.Conv2d(3, 32, 3, 2, 1), torch.nn.ReLU(),
+            torch.nn.Conv2d(32, 64, 3, 2, 1), torch.nn.ReLU(),
+            torch.nn.Conv2d(64, 128, 3, 2, 1), torch.nn.ReLU(),
+            torch.nn.AdaptiveAvgPool2d((1,1))
+        )
+        self.fc = torch.nn.Linear(128, embed_dim)
+    def forward(self, x):
+        x = self.cnn(x)
+        x = x.view(x.size(0), -1)
+        return self.fc(x)
+
+class TransformerDecoder(torch.nn.Module):
+    def __init__(self, vocab_size, embed_dim=128, nhead=4, num_layers=2, max_len=40):
+        super().__init__()
+        self.embed = torch.nn.Embedding(vocab_size, embed_dim)
+        decoder_layer = torch.nn.TransformerDecoderLayer(d_model=embed_dim, nhead=nhead)
+        self.decoder = torch.nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
+        self.fc_out = torch.nn.Linear(embed_dim, vocab_size)
+        self.pos_embed = torch.nn.Embedding(max_len, embed_dim)
+
+    def forward(self, tgt, memory):
+        positions = torch.arange(0, tgt.shape[1], device=tgt.device).unsqueeze(0)
+        tgt_emb = self.embed(tgt) + self.pos_embed(positions)
+        memory = memory.unsqueeze(0)
+        out = self.decoder(tgt_emb.transpose(0,1), memory)
+        return self.fc_out(out.transpose(0,1))
+
+class ImageCaptionModel(torch.nn.Module):
+    def __init__(self, vocab_size, embed_dim=128):
+        super().__init__()
+        self.encoder = CNNEncoder(embed_dim)
+        self.decoder = TransformerDecoder(vocab_size, embed_dim)
+    def forward(self, images, captions):
+        feats = self.encoder(images)
+        return self.decoder(captions, feats)
+
+# ============================================================
+# Load the tokenizer and model manually
+# ============================================================
+
+# Load config
+with open("radiology_caption_model/config.json", "r") as f:
+    config = json.load(f)
+
+# Load tokenizer - Corrected to use ByteLevelBPETokenizer with both files
+tokenizer = ByteLevelBPETokenizer("radiology_caption_model/vocab.json", "radiology_caption_model/merges.txt")
+
+# Instantiate the model with config parameters
+model = ImageCaptionModel(
+    vocab_size=config["vocab_size"],
+    embed_dim=config["embed_dim"]
+)
+
+# Load the model weights
+model.load_state_dict(torch.load("radiology_caption_model/pytorch_model.bin", map_location=torch.device('cpu')))
+model.eval() # Set model to evaluation mode
+
+# Define image transformations
+image_size = 128 # Must match training image size
+img_transforms = transforms.Compose([
+    transforms.Resize((image_size, image_size)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485,0.456,0.406],
+                         std=[0.229,0.224,0.225]),
+])
 
 # Function to generate caption
 def generate_caption(image):
     # Preprocess image
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Generate features
+    img_tensor = img_transforms(image).unsqueeze(0) # Add batch dimension
+
+    # Generate caption
     with torch.no_grad():
-        outputs = model.generate(**inputs)
-    
-    # Decode the output tokens
-    caption = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return caption
+        # Get image features
+        image_features = model.encoder(img_tensor)
+
+        # Start caption generation with BOS token
+        # (We assume BOS token ID is 2 from tokenizer training in cell 1)
+        # (Padding token ID is 0)
+        caption_tokens = [tokenizer.token_to_id("[BOS]")]
+        max_len = config["max_len"] if "max_len" in config else 40 # Use max_len from config, fallback to 40
+
+        for _ in range(max_len - 1): # -1 because BOS is already there
+            input_tokens = torch.tensor(caption_tokens).unsqueeze(0) # Add batch dimension
+            output = model.decoder(input_tokens, image_features)
+            last_token_logits = output[0, -1, :]
+            predicted_token_id = torch.argmax(last_token_logits).item()
+
+            caption_tokens.append(predicted_token_id)
+
+            # Stop if EOS token is generated
+            if predicted_token_id == tokenizer.token_to_id("[EOS]"):
+                break
+
+    # Decode the output tokens, excluding BOS and EOS (if present)
+    decoded_caption = tokenizer.decode(caption_tokens[1:-1] if caption_tokens[-1] == tokenizer.token_to_id("[EOS]") else caption_tokens[1:])
+    return decoded_caption
 
 # Create Gradio interface
 interface = gr.Interface(