feat: "first commit"

modality_connector.py

@@ -11,22 +11,27 @@ 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__()
-        self.processor = UnifiedImageProcessor(config)
-        self.model = ViTModel(ViTConfig.from_json_file(os.path.join(self.config_path, 'config.json')))
+
+        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))
+
         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,
@@ -68,16 +73,16 @@ class VisionEncoder(nn.Module):
 
 
 class UnifiedImageProcessor(nn.Module):
-    config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'vit_config')
-    def __init__(self, config: AuroraConfig):
+    def __init__(self, config: AuroraConfig, vit_config_path: str):
         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)
 
-        # Define resizer for pseudo-images (matches real image target size)
-        self.pseudo_resizer = Resize((self.target_size, self.target_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"]
+        self.pseudo_resizer = Resize((self.target_size, self.target_size))
         self.token_len = config.token_len
 
     def process_real_image(self, images):
@@ -107,7 +112,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)
 
@@ -124,20 +129,24 @@ 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__()
-        self.model = BertModel(BertConfig.from_json_file(os.path.join(self.config_path, 'config.json')))
+
+        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))
+
         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

@@ -9,7 +9,23 @@ from transformers.utils import ModelOutput
 
 
 class TSGenerationMixin(GenerationMixin):
-    tokenizer = BertTokenizer.from_pretrained(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'bert_config'), local_files_only=True)
+    _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")
+
+            if not os.path.isdir(tokenizer_dir):
+                raise FileNotFoundError(
+                    f"BERT tokenizer directory not found: {tokenizer_dir}"
+                )
+
+            self._tokenizer = BertTokenizer.from_pretrained(
+                tokenizer_dir,
+                local_files_only=True
+            )
+        return self._tokenizer
 
     @torch.no_grad()
     def generate(
@@ -95,7 +111,8 @@ class TSGenerationMixin(GenerationMixin):
         }
 
     def _tokenize(self, texts, max_length):
-        return self.tokenizer(
+        tokenizer = self._get_tokenizer()
+        return tokenizer(
             texts,
             padding='max_length',
             truncation=True,