feat: "first commit"

modality_connector.py

@@ -11,28 +11,22 @@ from .configuration_aurora import AuroraConfig
 
 
 class VisionEncoder(nn.Module):
+    config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'vit_config')
     def __init__(self, config: AuroraConfig):
         super().__init__()
-
-        base_dir = os.path.dirname(os.path.abspath(__file__))
-        self.config_path = os.path.join(base_dir, "vit_config")
-
-        self.processor = UnifiedImageProcessor(config, self.config_path)
-
-        vit_config_file = os.path.join(self.config_path, "config.json")
-
-        self.model = ViTModel(ViTConfig.from_json_file(vit_config_file))
-
+        self.processor = UnifiedImageProcessor(config)
+        self.model = ViTModel(ViTConfig.from_json_file(os.path.join(self.config_path, 'config.json')))
         for param in self.model.parameters():
             param.requires_grad = False
-
         self.hidden_size = self.model.config.hidden_size
         self.output_dim = config.hidden_size
         self.num_distill = config.num_distill
 
         self.projection = nn.Linear(self.hidden_size, self.output_dim)
+
         self.target_vision_tokens = nn.Parameter(torch.randn(self.num_distill, self.output_dim))
 
+        # Cross-attention layer
         self.cross_vision = nn.TransformerDecoder(
             nn.TransformerDecoderLayer(
                 d_model=config.hidden_size,
@@ -74,16 +68,16 @@ class VisionEncoder(nn.Module):
 
 
 class UnifiedImageProcessor(nn.Module):
-    def __init__(self, config: AuroraConfig, vit_config_path: str):
+    config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'vit_config')
+    def __init__(self, config: AuroraConfig):
         super().__init__()
+        # Load ViT preprocessor to get pretrained normalization parameters and target size
+        self.vit_processor = ViTImageProcessor.from_json_file(os.path.join(self.config_path, 'preprocessor_config.json'))
+        self.target_size = self.vit_processor.size["height"]  # e.g., 224 (default ViT input size)
 
-        self.config_path = vit_config_path
-
-        processor_file = os.path.join(self.config_path, "preprocessor_config.json")
-        self.vit_processor = ViTImageProcessor.from_json_file(processor_file)
-
-        self.target_size = self.vit_processor.size["height"]
+        # Define resizer for pseudo-images (matches real image target size)
         self.pseudo_resizer = Resize((self.target_size, self.target_size))
+
         self.token_len = config.token_len
 
     def process_real_image(self, images):
@@ -113,7 +107,7 @@ class UnifiedImageProcessor(nn.Module):
             period = input_length
 
         padding_length = (period - (input_length %
-                                    period)) % period
+                                            period)) % period
         x_pad = F.pad(x, (padding_length, 0))
         x_2d = einops.rearrange(x_pad, 'b (p f) -> b 1 f p', f=period)
 
@@ -130,25 +124,20 @@ class UnifiedImageProcessor(nn.Module):
 
 
 class TextEncoder(nn.Module):
+    config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'bert_config')
     def __init__(self, config: AuroraConfig):
         super().__init__()
-
-        base_dir = os.path.dirname(os.path.abspath(__file__))
-        self.config_path = os.path.join(base_dir, "bert_config")
-
-        bert_config_file = os.path.join(self.config_path, "config.json")
-
-        self.model = BertModel(BertConfig.from_json_file(bert_config_file))
-
+        self.model = BertModel(BertConfig.from_json_file(os.path.join(self.config_path, 'config.json')))
         for param in self.model.parameters():
             param.requires_grad = False
-
         self.hidden_size = self.model.config.hidden_size
         self.output_dim = config.hidden_size
         self.num_distill = config.num_distill
         self.max_length = 125
 
         self.projection = nn.Linear(self.hidden_size, self.output_dim)
+
+        # Define learnable target tokens (shape: [num_distill_tokens, hidden_size])
         self.target_text_tokens = nn.Parameter(torch.randn(self.num_distill, self.output_dim))
 
         self.cross_text = nn.TransformerDecoder(

ts_generation_mixin.py

@@ -7,19 +7,9 @@ from transformers import GenerationMixin, LogitsProcessorList, StoppingCriteriaL
 from transformers.generation.utils import GenerationConfig, GenerateOutput
 from transformers.utils import ModelOutput
 
-class TSGenerationMixin(GenerationMixin):
-    _tokenizer = None
-
-    def _get_tokenizer(self):
-        if self._tokenizer is None:
-            base_dir = os.path.dirname(os.path.abspath(__file__))
-            tokenizer_dir = os.path.join(base_dir, "bert_config")
 
-            self._tokenizer = BertTokenizer.from_pretrained(
-                tokenizer_dir,
-                local_files_only=True
-            )
-        return self._tokenizer
+class TSGenerationMixin(GenerationMixin):
+    tokenizer = BertTokenizer.from_pretrained(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'bert_config'), local_files_only=True)
 
     @torch.no_grad()
     def generate(
@@ -105,8 +95,7 @@ class TSGenerationMixin(GenerationMixin):
         }
 
     def _tokenize(self, texts, max_length):
-        tokenizer = self._get_tokenizer()
-        return tokenizer(
+        return self.tokenizer(
             texts,
             padding='max_length',
             truncation=True,