| 60 | |
| 61 | class VLM(nn.Module, GenerationMixin, ABC): |
| 62 | def __init__( |
| 63 | self, |
| 64 | model_family: str, |
| 65 | model_id: str, |
| 66 | vision_backbone: DinoSigLIPViTBackbone, |
| 67 | llm_backbone: LLaMa2LLMBackbone, |
| 68 | enable_mixed_precision_training: bool = True, |
| 69 | ) -> None: |
| 70 | super().__init__() |
| 71 | self.model_family, self.model_id = model_family, model_id |
| 72 | self.vision_backbone, self.llm_backbone = vision_backbone, llm_backbone |
| 73 | self.enable_mixed_precision_training = enable_mixed_precision_training |
| 74 | |
| 75 | # Instance Attributes for a generic VLM |
| 76 | self.all_module_keys, self.trainable_module_keys = None, None |
| 77 | |
| 78 | # === GenerationMixin Expected Attributes =>> *DO NOT MODIFY* === |
| 79 | self.generation_config = self.llm_backbone.llm.generation_config |
| 80 | self.main_input_name = "input_ids" |
| 81 | |
| 82 | @property |
| 83 | def device(self) -> torch.device: |