import torch

from torch import nn

from biovil_t.pretrained import get_biovil_t_image_encoder

class ChexpertClassifier(nn.Module):

    def __init__(self, num_classes):

        super().__init__()

        self.biovil_encoder = get_biovil_t_image_encoder()

        self.fc1 = nn.Linear(128 * 4 * 4, 512)

        self.fc2 = nn.Linear(512, num_classes)

    def forward(self, x):

        x = self.biovil_encoder(x).projected_patch_embeddings

        x = torch.nn.functional.avg_pool2d(x, 4)

        x = x.view(x.shape[0], -1)  # Flatten the tensor

        x = torch.relu(self.fc1(x))

        # x = self.biovil_encoder(x).img_embedding

        return self.fc2(x)