Add multimodal model support

.github/workflows/largemodel_unit_test_CI.yml

@@ -34,7 +34,7 @@ jobs:
           mode: start
           github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
           ec2-image-id: ${{ secrets.CUDA_EC2_IMAGE_ID }}
-          ec2-instance-type: g4dn.xlarge
+          ec2-instance-type: g4dn.2xlarge
           subnet-id: ${{ secrets.CUDA_SUBNET_ID }}
           security-group-id: ${{ secrets.CUDA_SECURITY_GROUP_ID }}
 
@@ -62,6 +62,13 @@ jobs:
           repository: marqo-ai/marqo-base
           path: marqo-base
 
+      - name: Install FFmpeg and libmagic
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y ffmpeg libmagic1
+          ffmpeg -version  # Verify installation
+          file --version   # Verify libmagic installation and version
+
       - name: Install dependencies
         run: |
           pip install -r marqo-base/requirements.txt

.github/workflows/unit_test_200gb_CI.yml

@@ -65,6 +65,12 @@ jobs:
           repository: marqo-ai/marqo-base
           path: marqo-base
 
+      - name: Install FFmpeg and libmagic
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y ffmpeg libmagic1
+          ffmpeg -version  # Verify installation
+          file --version   # Verify libmagic installation and version
       - name: Install dependencies
         run: |
           pip install -r marqo-base/requirements.txt

.gitignore

@@ -149,6 +149,4 @@ dump.rdb
 .DS_Store
 
 # Tester app for unit tests
-scripts/vespa_local/vespa_tester_app.zip
-
-src/marqo/tensor_search/cache_dir/*
+scripts/vespa_local/vespa_tester_app.zip

requirements.dev.txt

@@ -61,3 +61,13 @@ httpx==0.25.0
 # test requirements
 pyvespa==0.37.1
 pytest==7.4.3
+
+
+# LanguageBind
+eva-decord==0.6.1
+einops==0.6.1
+pytorchvideo==0.1.5
+torchaudio==0.12.1
+SoundFile==0.12.1
+python-magic==0.4.27
+ffmpeg-python==0.2.0

requirements.txt

@@ -1,14 +1,3 @@
-# This is a reduced set of requirements used for Marqo's tensor search module
-# If you are running this on your local machine you will need to install additional
-# requirements specified here: https://github.com/marqo-ai/marqo-base/blob/main/requirements.txt
-requests==2.28.1
-anyio==3.7.1
-fastapi==0.86.0
-uvicorn[standard]
-fastapi-utils==0.2.1
-jsonschema==4.17.1
-typing-extensions==4.5.0
-urllib3==1.26.0
 pydantic==1.10.11
 httpx==0.25.0
 semver==3.0.2
@@ -18,7 +7,4 @@ cachetools==5.3.1
 pynvml==11.5.0 # For cuda utilization
 readerwriterlock==1.0.9
 kazoo==2.10.0
-pycurl==7.45.3
-certifi==2019.11.28
-transformers==4.41.2
-optimum==1.20.0
+pycurl==7.45.3

src/marqo/api/configs.py

@@ -29,6 +29,8 @@ def default_env_vars() -> dict:
         EnvVars.MARQO_THREAD_EXPIRY_TIME: 1800,  # 30 minutes
         EnvVars.MARQO_ENABLE_THROTTLING: "TRUE",
         EnvVars.MARQO_LOG_LEVEL: "info",
+        EnvVars.MARQO_MEDIA_DOWNLOAD_THREAD_COUNT_PER_REQUEST: 5,
+        EnvVars.MARQO_IMAGE_DOWNLOAD_THREAD_COUNT_PER_REQUEST: 20,
         # This env variable is set to "info" by default in run_marqo.sh, which overrides this value
         EnvVars.MARQO_MAX_CPU_MODEL_MEMORY: 4,
         EnvVars.MARQO_MAX_CUDA_MODEL_MEMORY: 4,  # For multi-GPU, this is the max memory for each GPU.

src/marqo/core/models/marqo_index.py

@@ -38,6 +38,8 @@ class FieldType(str, Enum):
     ArrayFloat = 'array<float>'
     ArrayDouble = 'array<double>'
     ImagePointer = 'image_pointer'
+    VideoPointer = 'video_pointer'
+    AudioPointer = 'audio_pointer'
     MultimodalCombination = 'multimodal_combination'
     CustomVector = "custom_vector"
     MapInt = 'map<text, int>'
@@ -117,6 +119,13 @@ class TextPreProcessing(ImmutableStrictBaseModel):
     split_overlap: int = pydantic.Field(ge=0, alias='splitOverlap')
     split_method: TextSplitMethod = pydantic.Field(alias='splitMethod')
 
+class VideoPreProcessing(ImmutableStrictBaseModel):
+    split_length: int = pydantic.Field(gt=0, alias='splitLength')
+    split_overlap: int = pydantic.Field(ge=0, alias='splitOverlap')
+
+class AudioPreProcessing(ImmutableStrictBaseModel):
+    split_length: int = pydantic.Field(gt=0, alias='splitLength')
+    split_overlap: int = pydantic.Field(ge=0, alias='splitOverlap')
 
 class ImagePreProcessing(ImmutableStrictBaseModel):
     patch_method: Optional[PatchMethod] = pydantic.Field(alias='patchMethod')
@@ -250,6 +259,8 @@ class MarqoIndex(ImmutableBaseModel, ABC):
     normalize_embeddings: bool
     text_preprocessing: TextPreProcessing
     image_preprocessing: ImagePreProcessing
+    video_preprocessing: Optional[VideoPreProcessing] = None
+    audio_preprocessing: Optional[AudioPreProcessing] = None
     distance_metric: DistanceMetric
     vector_numeric_type: VectorNumericType
     hnsw_config: HnswConfig
@@ -324,6 +335,7 @@ def _cache_or_get(self, key: str, func):
 class UnstructuredMarqoIndex(MarqoIndex):
     type = IndexType.Unstructured
     treat_urls_and_pointers_as_images: bool
+    treat_urls_and_pointers_as_media: Optional[bool] = None
     filter_string_max_length: int
 
     def __init__(self, **data):

src/marqo/core/models/marqo_index_request.py

@@ -24,6 +24,8 @@ class MarqoIndexRequest(ImmutableStrictBaseModel, ABC):
     normalize_embeddings: bool
     text_preprocessing: marqo_index.TextPreProcessing
     image_preprocessing: marqo_index.ImagePreProcessing
+    video_preprocessing: marqo_index.VideoPreProcessing
+    audio_preprocessing: marqo_index.AudioPreProcessing
     distance_metric: marqo_index.DistanceMetric
     vector_numeric_type: marqo_index.VectorNumericType
     hnsw_config: marqo_index.HnswConfig
@@ -39,9 +41,9 @@ def validate_name(cls, name):
 
 class UnstructuredMarqoIndexRequest(MarqoIndexRequest):
     treat_urls_and_pointers_as_images: bool
+    treat_urls_and_pointers_as_media: bool
     filter_string_max_length: int
 
-
 class FieldRequest(StrictBaseModel):
     name: str
     type: marqo_index.FieldType

src/marqo/core/search/recommender.py

@@ -43,7 +43,7 @@ def recommend(self,
         Args:
             index_name: Name of the index to search
             documents: A list of document IDs or a dictionary where the keys are document IDs and the values are weights
-            tensor_fields: List of tensor fields to use for recommendation
+            tensor_fields: List of tensor fields to use for recommendation (can include text, image, audio, and video fields)
             interpolation_method: Interpolation method to use for combining vectors
             exclude_input_documents: Whether to exclude the input documents from the search results
             result_count: Number of results to return

src/marqo/core/structured_vespa_index/structured_vespa_index.py

@@ -30,6 +30,8 @@ class StructuredVespaIndex(VespaIndex):
         FieldType.ArrayLong: (list, int),
         FieldType.ArrayDouble: (list, (float, int)),
         FieldType.ImagePointer: str,
+        FieldType.VideoPointer: str,
+        FieldType.AudioPointer: str,
         FieldType.MultimodalCombination: dict,
         FieldType.CustomVector: str,
         FieldType.MapInt: (dict, int),

src/marqo/core/structured_vespa_index/structured_vespa_schema.py

@@ -22,6 +22,8 @@ class StructuredVespaSchema(VespaSchema):
         FieldType.ArrayFloat: 'array<float>',
         FieldType.ArrayDouble: 'array<double>',
         FieldType.ImagePointer: 'string',
+        FieldType.VideoPointer: 'string',
+        FieldType.AudioPointer: 'string',
         FieldType.MultimodalCombination: 'map<string, float>',
         FieldType.CustomVector: 'string',  # Custom Vector "content" is stored as string in backend.
         FieldType.MapInt: 'map<string, int>',
@@ -165,6 +167,8 @@ def _generate_document_section(self, schema_name: str) -> (List[str], Structured
             normalize_embeddings=self._index_request.normalize_embeddings,
             text_preprocessing=self._index_request.text_preprocessing,
             image_preprocessing=self._index_request.image_preprocessing,
+            video_preprocessing=self._index_request.video_preprocessing,
+            audio_preprocessing=self._index_request.audio_preprocessing,
             distance_metric=self._index_request.distance_metric,
             vector_numeric_type=self._index_request.vector_numeric_type,
             hnsw_config=self._index_request.hnsw_config,

src/marqo/core/unstructured_vespa_index/unstructured_vespa_schema.py

@@ -54,13 +54,16 @@ def _generate_unstructured_marqo_index(self, schema_name: str) -> UnstructuredMa
             normalize_embeddings=self._index_request.normalize_embeddings,
             text_preprocessing=self._index_request.text_preprocessing,
             image_preprocessing=self._index_request.image_preprocessing,
+            video_preprocessing=self._index_request.video_preprocessing,
+            audio_preprocessing=self._index_request.audio_preprocessing,
             distance_metric=self._index_request.distance_metric,
             vector_numeric_type=self._index_request.vector_numeric_type,
             hnsw_config=self._index_request.hnsw_config,
             marqo_version=self._index_request.marqo_version,
             created_at=self._index_request.created_at,
             updated_at=self._index_request.updated_at,
             treat_urls_and_pointers_as_images=self._index_request.treat_urls_and_pointers_as_images,
+            treat_urls_and_pointers_as_media=self._index_request.treat_urls_and_pointers_as_media,
             filter_string_max_length=self._index_request.filter_string_max_length,
         )
 

src/marqo/s2_inference/clip_utils.py

@@ -896,4 +896,4 @@ def get_multilingual_clip_properties() -> Dict:
                 "type": "multilingual_clip",
             }
     }
-    return MULTILINGUAL_CLIP_PROPERTIES
+    return MULTILINGUAL_CLIP_PROPERTIES

src/marqo/s2_inference/configs.py

@@ -14,13 +14,16 @@ class ModelCache:
     # The hf_cache_path is managed by the hf_hub_download function
     hf_cache_path = os.getenv('HF_SAVE_PATH', f'{utils.get_marqo_root_from_env()}/cache/hf/')
 
+    languagebind_cache_path = os.getenv('LANGUAGEBIND_CACHE_PATH', f'{utils.get_marqo_root_from_env()}/cache/languagebind/')
+
     @classmethod
     def get_all_cache_paths(cls):
         return [
             cls.onnx_cache_path,
             cls.torch_cache_path,
             cls.clip_cache_path,
-            cls.hf_cache_path
+            cls.hf_cache_path,
+            cls.languagebind_cache_path
         ]
 
 class BaseTransformerModels:

src/marqo/s2_inference/errors.py

@@ -62,4 +62,8 @@ class BatchInferenceSizeNotMatchError(S2InferenceError):
 
 
 class ImageDownloadError(S2InferenceError):
+    pass
+
+
+class UnsupportedModalityError(S2InferenceError):
     pass

src/marqo/s2_inference/languagebind/__init__.py

@@ -0,0 +1,69 @@
+from torch import nn
+
+from .audio.configuration_audio import LanguageBindAudioConfig
+from .audio.modeling_audio import LanguageBindAudio
+from .audio.processing_audio import LanguageBindAudioProcessor
+from .audio.tokenization_audio import LanguageBindAudioTokenizer
+from .image.configuration_image import LanguageBindImageConfig
+from .image.modeling_image import LanguageBindImage
+from .image.processing_image import LanguageBindImageProcessor
+from .image.tokenization_image import LanguageBindImageTokenizer
+from .video.configuration_video import LanguageBindVideoConfig
+from .video.modeling_video import LanguageBindVideo
+from .video.processing_video import LanguageBindVideoProcessor
+from .video.tokenization_video import LanguageBindVideoTokenizer
+
+config_dict = {
+    'image': LanguageBindImageConfig,
+    'video': LanguageBindVideoConfig,
+    'audio': LanguageBindAudioConfig
+}
+model_dict = {
+    'image': LanguageBindImage,
+    'video': LanguageBindVideo,
+    'audio': LanguageBindAudio
+}
+transform_dict = {
+    'video': LanguageBindVideoProcessor,
+    'audio': LanguageBindAudioProcessor,
+    'image': LanguageBindImageProcessor,
+}
+
+
+class LanguageBind(nn.Module):
+    def __init__(self, clip_type, use_temp=True, cache_dir='./cache_dir'):
+        super(LanguageBind, self).__init__()
+        self.use_temp = use_temp
+        self.modality_encoder = {}
+        self.modality_proj = {}
+        self.modality_scale = {}
+        self.modality_config = {}
+        for k, v in clip_type.items():
+            pretrained_ckpt = f'LanguageBind/{v}'
+            model = model_dict[k].from_pretrained(pretrained_ckpt, cache_dir=cache_dir)
+            self.modality_encoder[k] = model.vision_model
+            self.modality_proj[k] = model.visual_projection
+            self.modality_scale[k] = model.logit_scale
+            self.modality_config[k] = model.config
+        self.modality_encoder['language'] = model.text_model
+        self.modality_proj['language'] = model.text_projection
+
+        self.modality_encoder = nn.ModuleDict(self.modality_encoder)
+        self.modality_proj = nn.ModuleDict(self.modality_proj)
+
+    def forward(self, inputs):
+        outputs = {}
+        for key, value in inputs.items():
+            value = self.modality_encoder[key](value)[1]
+            value = self.modality_proj[key](value)
+            value = value / value.norm(p=2, dim=-1, keepdim=True)
+            if self.use_temp:
+                if key != 'language':
+                    value = value * self.modality_scale[key].exp()
+            outputs[key] = value
+        return outputs
+
+
+def to_device(x, device):
+    out_dict = {k: v.to(device) for k, v in x.items()}
+    return out_dict