feat: "first commit"

modality_connector.py

@@ -8,6 +8,7 @@ from torchvision.transforms import Resize
 from transformers import ViTImageProcessor, ViTModel, BertModel, ViTConfig, BertConfig
 
 from .configuration_aurora import AuroraConfig
+from .util_functions import resolve_subdir
 
 
 class VisionEncoder(nn.Module):
@@ -20,6 +21,10 @@ class VisionEncoder(nn.Module):
         self.processor = UnifiedImageProcessor(config, self.config_path)
 
         vit_config_file = os.path.join(self.config_path, "config.json")
+
+        if not os.path.exists(vit_config_file):
+            resolve_subdir(repo_id="DecisionIntelligence/Aurora", subdir="vit_config", file_name="config.json")
+
         self.model = ViTModel(ViTConfig.from_json_file(vit_config_file))
 
         for param in self.model.parameters():
@@ -79,6 +84,9 @@ class UnifiedImageProcessor(nn.Module):
         self.config_path = vit_config_path
 
         processor_file = os.path.join(self.config_path, "preprocessor_config.json")
+        if not os.path.exists(processor_file):
+            resolve_subdir(repo_id="DecisionIntelligence/Aurora", subdir="vit_config", file_name="preprocessor_config.json")
+
         self.vit_processor = ViTImageProcessor.from_json_file(processor_file)
 
         self.target_size = self.vit_processor.size["height"]
@@ -136,6 +144,10 @@ class TextEncoder(nn.Module):
         self.config_path = os.path.join(base_dir, "bert_config")
 
         bert_config_file = os.path.join(self.config_path, "config.json")
+
+        if not os.path.exists(bert_config_file):
+            resolve_subdir(repo_id="DecisionIntelligence/Aurora", subdir="bert_config", file_name="config.json")
+
         self.model = BertModel(BertConfig.from_json_file(bert_config_file))
 
         for param in self.model.parameters():

ts_generation_mixin.py

@@ -6,7 +6,7 @@ from transformers import BertTokenizer
 from transformers import GenerationMixin, LogitsProcessorList, StoppingCriteriaList
 from transformers.generation.utils import GenerationConfig, GenerateOutput
 from transformers.utils import ModelOutput
-
+from .util_functions import resolve_subdir
 
 class TSGenerationMixin(GenerationMixin):
     _tokenizer = None
@@ -17,9 +17,7 @@ class TSGenerationMixin(GenerationMixin):
             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}"
-                )
+                resolve_subdir(repo_id="DecisionIntelligence/Aurora", subdir="bert_config", file_name="config.json")
 
             self._tokenizer = BertTokenizer.from_pretrained(
                 tokenizer_dir,

util_functions.py

@@ -1,9 +1,18 @@
+import os
 from typing import Tuple
 
 import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+
+
+def resolve_subdir(repo_id: str, subdir: str, file_name: str) -> str:
+    hf_hub_download(
+        repo_id=repo_id,
+        filename=f"{subdir}/{file_name}"
+    )
 
 
 def resize(x_tensor, new_shape):