| --- |
| license: mit |
| --- |
| |
| # BLIPNet Model |
|
|
| This is the structure of the BLIPNet model. You can load the model with this structure, or you can create a bigger model for your specific task. |
|
|
| ## Model Structure |
|
|
| ```python |
| import torch |
| import torch.nn as nn |
| from transformers import BlipForConditionalGeneration |
| |
| class BLIPNet(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| # Generation Model |
| self.model = BlipForConditionalGeneration.from_pretrained("Salesforceblip-image-captioning-base", cache_dir="model") |
| # Same with https://huggingface.co/uf-aice-lab/BLIP-Math |
| self.ebd_dim = 443136 |
| |
| # Classification Model |
| fc_dim = 64 # You can choose a higher number for better performance, for example, 1024. |
| self.head = nn.Sequential( |
| nn.Linear(self.ebd_dim, fc_dim), |
| nn.ReLU(), |
| ) |
| self.output1= nn.Linear(fc_dim, 5) # 5 classes |
| |
| def forward(self, pixel_values, input_ids): |
| outputs = self.model(input_ids=input_ids, pixel_values=pixel_values, labels=input_ids) |
| image_text_embeds = self.model.vision_model(pixel_values, return_dict=True).last_hidden_state |
| image_text_embeds = self.head(image_text_embeds.view(-1, self.ebd_dim)) |
| |
| # A classification model is based on embeddings from a generative model to leverage BLIP's powerful image-text encoding capabilities. |
| logits = self.output1(image_text_embeds) |
| |
| # generated text, probabilities of classification |
| return outputs, logits |
| |
| model = BLIPNet() |
| model.load_state_dict(torch.load("BLILP_Generation_Classification.bin"), strict=False) |
| |
| You need to input the sample in the same way as shown in the example provided at: https://huggingface.co/uf-aice-lab/BLIP-Math |
| Then you can get the generated text and classification score simultaneously. |