newer model

model.py

@@ -0,0 +1,85 @@
+
+class Model(PreTrainedModel):
+    config_class = VLMConfig
+    
+    def __init__(self, config: VLMConfig, image_model, language_model, num_projections: int, tokenizer, prepend_text: str, image_tokens:int):
+        super().__init__(config)
+        self.image_model = image_model
+        self.language_model = language_model
+        self.projector = nn.Sequential(
+            *projection_layers(image_model.num_features, language_model.config.hidden_size, num_projections)
+        )
+        
+        self.tokenizer = tokenizer
+        self.eos_token = tokenizer.eos_token
+        self.prepend_text = prepend_text
+        
+        self.image_tokens = image_tokens
+        
+        input_ids = tokenizer(prepend_text, return_tensors="pt").input_ids
+        eos_token_index = (input_ids[0] == tokenizer.eos_token_id).nonzero(as_tuple=True)[0].item()
+        text_embeddings = self.language_model.get_input_embeddings()(input_ids).detach()
+        self.prepend_embeddings = text_embeddings[:, :eos_token_index]
+        self.postpend_embeddings = text_embeddings[:, eos_token_index:]
+        self.attention_mask = torch.ones(1, text_embeddings.shape[1] + image_tokens)
+        self.labels = torch.full((1, self.attention_mask.shape[1]), LABEL_MASK)
+        
+    def project_image_features(self, images: torch.Tensor):
+        image_features = self.image_model.forward_features(images)
+        image_features = einops.rearrange(image_features, "bs dim w h -> bs (w h) dim")
+        encoder_outputs = self.projector(image_features)
+        return encoder_outputs
+        
+    def forward(self, images: torch.Tensor, tokenized_captions: dict[str, torch.Tensor]):
+        image_outputs = self.project_image_features(images)
+        caption_embeddings = self.language_model.get_input_embeddings()(tokenized_captions.input_ids).detach()
+        device = images.device
+        embeddings = torch.cat(
+            [
+                self.prepend_embeddings.to(device).expand(len(images), -1, -1),
+                image_outputs,
+                self.postpend_embeddings.to(device).expand(len(images), -1, -1),
+                caption_embeddings,
+            ],
+            dim=1,
+        )
+        attention_mask = torch.cat(
+            [
+                self.attention_mask.to(device).expand(len(images), -1), 
+                tokenized_captions.attention_mask
+            ], 
+            dim=1
+        )
+        labels = torch.cat(
+            [
+                self.labels.to(device).expand(len(images), -1), 
+                tokenized_captions.input_ids.clone()
+            ],
+            dim=1,
+        )
+        labels[attention_mask == 0] = LABEL_MASK
+        
+        return self.language_model(
+            inputs_embeds=embeddings,
+            attention_mask=attention_mask,
+            labels=labels,
+        )
+    
+    def generate(self, images: torch.Tensor, generator_kwargs: dict[str, Union[int, float]]):
+        image_outputs = self.project_image_features(images)
+        device = images.device
+        embeddings = torch.cat(
+            [
+                self.prepend_embeddings.to(device).expand(len(images), -1, -1),
+                image_outputs,
+                self.postpend_embeddings.to(device).expand(len(images), -1, -1),
+            ],
+            dim=1,
+        )
+        attention_mask = self.attention_mask.to(device).expand(len(images), -1)
+        return self.language_model.generate(
+            inputs_embeds=embeddings,
+            attention_mask=attention_mask,
+            eos_token_id=self.tokenizer.eos_token_id,
+            **generator_kwargs
+        )