Initial commit

src/sonicverse/configs/tasks.json

@@ -0,0 +1,208 @@
+{
+    "backbone": {
+        "num_layers": 5,
+        "input_channels": 25,
+        "output_channels": 25,
+        "output_size": 4096,
+        "hidden_size": 4096,
+        "requires_grad": true
+    },
+    "task_heads": {
+        "lmm_projector": {
+            "num_layers": 3,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "input_size": 768,
+            "input_channels": 13,
+            "width":  40,
+            "weight": 1.0,
+            "model_type": "mlp",
+            "requires_grad": true,
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "use_backbone_output": false
+        },
+        "instrument_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 40,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.1,
+            "requires_grad": true,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "mood_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 56,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.1,
+            "requires_grad": true,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "genre_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 87,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.1,
+            "requires_grad": true,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "key_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 24,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.1,
+            "requires_grad": true,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "vocals_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 3,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.1,
+            "requires_grad": true,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        }
+    },
+    "task_projectors": {
+        "instrument_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 40,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "mood_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 56,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "genre_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 87,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "key_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 24,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "vocals_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 3,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "chords_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 216,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "beats_detection": {
+            "model_type": "mlp_conv_agg",
+            "num_layers": 3,
+            "input_channels": 2,
+            "input_size": 500,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": true,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": true,
+            "requires_grad": true
+        }
+    }
+}

src/sonicverse/configs/tasks_baseline.json

@@ -0,0 +1,20 @@
+{
+    "task_heads": {
+        "lmm_projector": {
+            "num_layers": 3,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "input_size": 768,
+            "input_channels": 13,
+            "width":  60,
+            "weight": 1.0,
+            "model_type": "mlp",
+            "requires_grad": true,
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": false,
+            "use_transpose": false
+        }
+    },
+    "task_projectors": {}
+}

src/sonicverse/configs/tasks_ft.json

@@ -0,0 +1,208 @@
+{
+    "backbone": {
+        "num_layers": 5,
+        "input_channels": 25,
+        "output_channels": 25,
+        "output_size": 4096,
+        "hidden_size": 4096,
+        "requires_grad": false
+    },
+    "task_heads": {
+        "lmm_projector": {
+            "num_layers": 3,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "input_size": 768,
+            "input_channels": 13,
+            "width":  40,
+            "weight": 1.0,
+            "model_type": "mlp",
+            "requires_grad": true,
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "use_backbone_output": false
+        },
+        "instrument_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 40,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.0,
+            "requires_grad": false,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "mood_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 56,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.0,
+            "requires_grad": false,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "genre_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 87,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.0,
+            "requires_grad": false,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "key_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 24,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.0,
+            "requires_grad": false,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        },
+        "vocals_detection": {
+            "model_type": "mlp",
+            "use_aggregator": true,
+            "use_time_average": true,
+            "use_sigmoid": true,
+            "use_transpose": false,
+            "num_layers": 2,
+            "input_size": 508,
+            "output_size": 3,
+            "hidden_size": 4096,
+            "width": 1,
+            "weight": 0.0,
+            "requires_grad": false,
+            "num_conv_layers": 4,
+            "output_channel": 1
+        }
+    },
+    "task_projectors": {
+        "instrument_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 40,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "mood_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 56,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "genre_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 87,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "key_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 24,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "vocals_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 3,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "chords_detection": {
+            "model_type": "mlp",
+            "num_layers": 3,
+            "input_channels": 0,
+            "input_size": 216,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": false,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": false,
+            "requires_grad": true
+        },
+        "beats_detection": {
+            "model_type": "mlp_conv_agg",
+            "num_layers": 3,
+            "input_channels": 2,
+            "input_size": 500,
+            "output_size": 4096,
+            "hidden_size": 4096,
+            "width": 4,
+            "use_aggregator": true,
+            "use_time_average": false,
+            "use_sigmoid": false,
+            "use_transpose": true,
+            "requires_grad": true
+        }
+    }
+}

src/sonicverse/configs/tasks_pt_weight.json

@@ -0,0 +1,10 @@
+{
+    "pretrained_paths": [
+        {
+            "path": "/experiments/music_extraction/mlp_shared_multi_task_trial_002/train_002_epoch_45_step_40.pth",
+            "components": ["backbone", "instrument_detection", "genre_detection", "mood_detection", "key_detection", "vocals_detection"],
+            "use_prefix": true,
+            "prefix": "audio_mert_lmm_projector"
+        }
+    ]
+}

src/sonicverse/configs/zero2.json

@@ -0,0 +1,23 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 2,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto"
+    }
+}

src/sonicverse/configs/zero3.json

@@ -0,0 +1,28 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 3,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto",
+        "stage3_prefetch_bucket_size": "auto",
+        "stage3_param_persistence_threshold": "auto",
+        "stage3_max_live_parameters": 1e9,
+        "stage3_max_reuse_distance": 1e9,
+        "stage3_gather_16bit_weights_on_model_save": true
+    }
+}

src/sonicverse/configs/zero3_offload.json

@@ -0,0 +1,56 @@
+{
+    "fp16": {
+      "enabled": "auto",
+      "loss_scale": 0,
+      "loss_scale_window": 1000,
+      "initial_scale_power": 16,
+      "hysteresis": 2,
+      "min_loss_scale": 1
+    },
+    "bf16": {
+      "enabled": "auto"
+    },
+    "optimizer": {
+      "type": "AdamW",
+      "params": {
+        "lr": "auto",
+        "betas": "auto",
+        "eps": "auto",
+        "weight_decay": "auto"
+      }
+    },
+    "scheduler": {
+      "type": "WarmupLR",
+      "params": {
+        "warmup_min_lr": "auto",
+        "warmup_max_lr": "auto",
+        "warmup_num_steps": "auto"
+      }
+    },
+    "zero_optimization": {
+      "stage": 3,
+      "offload_optimizer": {
+        "device": "cpu",
+        "pin_memory": true
+      },
+      "offload_param": {
+        "device": "cpu",
+        "pin_memory": true
+      },
+      "overlap_comm": true,
+      "contiguous_gradients": true,
+      "sub_group_size": 1e9,
+      "reduce_bucket_size": "auto",
+      "stage3_prefetch_bucket_size": "auto",
+      "stage3_param_persistence_threshold": "auto",
+      "stage3_max_live_parameters": 1e9,
+      "stage3_max_reuse_distance": 1e9,
+      "gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "steps_per_print": 1e5,
+    "wall_clock_breakdown": false
+  }

src/sonicverse/multi_token.egg-info/PKG-INFO

@@ -0,0 +1,6 @@
+Metadata-Version: 2.1
+Name: multi-token
+Version: 0.0.4
+Home-page: https://github.com/sshh12/multi_token
+Author: Shrivu Shankar
+License: Apache License 2.0

src/sonicverse/multi_token.egg-info/SOURCES.txt

@@ -0,0 +1,6 @@
+setup.py
+multi_token.egg-info/PKG-INFO
+multi_token.egg-info/SOURCES.txt
+multi_token.egg-info/dependency_links.txt
+multi_token.egg-info/requires.txt
+multi_token.egg-info/top_level.txt

src/sonicverse/multi_token.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

src/sonicverse/multi_token.egg-info/requires.txt

@@ -0,0 +1,8 @@
+transformers>=4.34.0
+accelerate>=0.21.0
+scipy>=1.11.3
+bitsandbytes>=0.41.0
+datasets>=2.14.5
+sentencepiece>=0.1.99
+peft>=0.4.0
+deepspeed==0.9.5

src/sonicverse/multi_token.egg-info/top_level.txt

@@ -0,0 +1 @@
+

src/sonicverse/multi_token/constants.py

@@ -0,0 +1,4 @@
+IGNORE_INDEX = -100
+
+ROLE_ASSISTANT = "assistant"
+ROLE_USER = "user"

src/sonicverse/multi_token/data_tools.py

@@ -0,0 +1,336 @@
+from typing import Dict, List, Any, Union, Optional
+from collections import Counter
+from functools import cache
+import contextlib
+import tempfile
+import shutil
+import random
+import subprocess
+import json
+import re
+import io
+import os
+
+import torch
+import requests
+import transformers
+import numpy as np
+from datasets import load_dataset, Dataset
+from PIL import Image
+
+from multi_token.constants import IGNORE_INDEX
+
+
+def encode_chat(
+    item: Dict,
+    tokenizer: transformers.PreTrainedTokenizer,
+    modalities: List["Modality"],
+) -> Dict:
+    messages = list(item["messages"])
+    chat_as_string = tokenizer.apply_chat_template(messages, tokenize=False)
+
+    token_to_modality = {m.token: m for m in modalities}
+    modality_token_counts = Counter()
+    instruct_pattern = r"(\[INST\][\s\S]*?\[\/INST\])"
+    pattern = "(" + "|".join(re.escape(m.token) for m in modalities) + ")"
+
+    chat_part = re.split(instruct_pattern, chat_as_string)
+    input_ids = []
+    labels = []
+    for part in chat_part:
+        if "[INST]" in part:
+            is_instruction = True
+        else:
+            is_instruction = False
+        for subpart in re.split(pattern, part):
+            if not subpart:
+                continue
+            if subpart in token_to_modality:
+                assert (
+                    is_instruction
+                ), "There should be no modality tokens outside of instructions"
+                m = token_to_modality[subpart]
+                modality_token_counts[m.name] += 1
+                input_ids.extend([m.token_idx] * m.token_width)
+                labels.extend([IGNORE_INDEX] * m.token_width)
+            elif is_instruction:
+                part_ids = tokenizer(subpart, add_special_tokens=False).input_ids
+                input_ids.extend(part_ids)
+                labels.extend([IGNORE_INDEX] * len(part_ids))
+            else:
+                part_ids = tokenizer(subpart, add_special_tokens=False).input_ids
+                input_ids.extend(part_ids)
+                labels.extend(part_ids)
+
+    input_ids = torch.tensor(input_ids, dtype=torch.long)
+    labels = torch.tensor(labels, dtype=torch.long)
+
+    data_dict = dict(
+        input_ids=input_ids,
+        labels=labels,
+    )
+    for m in modalities:
+        data_dict[m.name] = m.preprocess_rows([item])[0]
+    return data_dict
+
+def encode_chat_multitask(
+    item: Dict,
+    tokenizer: transformers.PreTrainedTokenizer,
+    modalities: List["Modality"],
+) -> Dict:
+    messages = list(item["messages"])
+    chat_as_string = tokenizer.apply_chat_template(messages, tokenize=False)
+
+    token_to_modality = {m.token: m for m in modalities}
+    modality_token_counts = Counter()
+    instruct_pattern = r"(\[INST\][\s\S]*?\[\/INST\])"
+    pattern = "(" + "|".join(re.escape(m.token) for m in modalities) + ")"
+
+    chat_part = re.split(instruct_pattern, chat_as_string)
+    input_ids = []
+    labels = []
+    labels.append([])
+    for part in chat_part:
+        if "[INST]" in part:
+            is_instruction = True
+        else:
+            is_instruction = False
+        for subpart in re.split(pattern, part):
+            if not subpart:
+                continue
+            if subpart in token_to_modality:
+                assert (
+                    is_instruction
+                ), "There should be no modality tokens outside of instructions"
+                m = token_to_modality[subpart]
+                modality_token_counts[m.name] += 1
+                input_ids.extend([m.token_idx] * m.token_width)
+                labels[0].extend([IGNORE_INDEX] * m.token_width)
+            elif is_instruction:
+                part_ids = tokenizer(subpart, add_special_tokens=False).input_ids
+                input_ids.extend(part_ids)
+                labels[0].extend([IGNORE_INDEX] * len(part_ids))
+            else:
+                part_ids = tokenizer(subpart, add_special_tokens=False).input_ids
+                input_ids.extend(part_ids)
+                labels[0].extend(part_ids)
+
+    input_ids = torch.tensor(input_ids, dtype=torch.long)
+    labels[0] = torch.tensor(labels[0], dtype=torch.long)
+
+    task_list = []
+    for m in modalities:
+        task_list += m.tasks["task_heads"].keys()
+    # labels[task_specs["task_id"]] = load_tensor(item[task_name][0])
+
+    for task_name in task_list:
+        if task_name != "lmm_projector":
+            labels.append(load_tensor(item[task_name][0]))
+
+    # labels = torch.tensor(labels, dtype=torch.long)
+
+    data_dict = dict(
+        input_ids=input_ids,
+        labels=labels,
+    )
+    for m in modalities:
+        data_dict[m.name] = m.preprocess_rows([item])[0]
+    return data_dict
+
+def load_tensor(path: str) -> np.ndarray:
+    return torch.tensor(np.load(path))
+
+
+def load_image(value: Any) -> Image.Image:
+    img = None
+    if isinstance(value, str):
+        if value.startswith("http://") or value.startswith("https://"):
+            response = requests.get(value)
+            img = Image.open(io.BytesIO(response.content))
+        elif os.path.exists(value):
+            img = Image.open(value)
+    elif isinstance(value, Image.Image):
+        img = value
+    if img is None:
+        raise ValueError(f"Could not load image from {value}")
+    img = img.convert("RGB")
+    return img
+
+
+@contextlib.contextmanager
+def with_local_files(fn_or_urls: List[Any]):
+    local_fns = []
+    fps = []
+    for fn_or_url in fn_or_urls:
+        if isinstance(fn_or_url, Image.Image):
+            fp = tempfile.NamedTemporaryFile(suffix=".png", mode="wb")
+            fn_or_url.convert("RGB").save(fp)
+            fps.append(fp)
+            local_fns.append(fp.name)
+        elif fn_or_url.startswith("http://") or fn_or_url.startswith("https://"):
+            suffix = os.path.splitext(fn_or_url)[-1]
+            with requests.get(fn_or_url, stream=True) as r:
+                fp = tempfile.NamedTemporaryFile(suffix=suffix, mode="wb")
+                shutil.copyfileobj(r.raw, fp)
+                fps.append(fp)
+                local_fns.append(fp.name)
+        else:
+            local_fns.append(fn_or_url)
+    try:
+        yield local_fns
+    finally:
+        for fp in fps:
+            fp.close()
+
+
+@cache
+def _get_dataset(dataset_args: str) -> Dataset:
+    return load_dataset(**json.loads(dataset_args))
+
+
+def get_dataset_cached(dataset_args: Dict) -> Dataset:
+    return _get_dataset(json.dumps(dataset_args))
+
+
+def load_audio_signal(input_: Union[Dict, str]) -> Dict:
+    from audiotools import AudioSignal
+
+    if isinstance(input_, dict) and "array" in input_:
+        array = input_["array"]
+    elif isinstance(input_, dict) and "dataset_args" in input_:
+        item = get_dataset_cached(input_["dataset_args"])[input_["idx"]]
+        array = item["audio"]["array"]
+    elif isinstance(input_, dict) and "path" in input_:
+        with with_local_files([input_["path"]]) as local_fns:
+            array =  AudioSignal(local_fns[0])
+    elif isinstance(input_, str):
+        with with_local_files([input_]) as local_fns:
+            array =  AudioSignal(local_fns[0])
+    else:
+        raise ValueError(f"Could not load audio from {input_}")
+
+    return {"array": list(array)}
+
+
+def load_audio(input_: Union[Dict, str], target_sampling_rate: int = None) -> Dict:
+    import soundfile as sf
+    import librosa
+
+    if isinstance(input_, dict) and "array" in input_ and "sampling_rate" in input_:
+        array = input_["array"]
+        sampling_rate = input_["sampling_rate"]
+    elif isinstance(input_, dict) and "dataset_args" in input_:
+        item = get_dataset_cached(input_["dataset_args"])[input_["idx"]]
+        array = item["audio"]["array"]
+        sampling_rate = item["audio"]["sampling_rate"]
+    elif isinstance(input_, dict) and "path" in input_:
+        with with_local_files([input_["path"]]) as local_fns:
+            array, sampling_rate = sf.read(local_fns[0])
+    elif isinstance(input_, str):
+        with with_local_files([input_]) as local_fns:
+            array, sampling_rate = sf.read(local_fns[0])
+    else:
+        raise ValueError(f"Could not load audio from {input_}")
+
+    if array.ndim == 2:
+        array = array.mean(axis=1)
+
+    if target_sampling_rate is not None and sampling_rate != target_sampling_rate:
+        array = librosa.resample(
+            array, orig_sr=sampling_rate, target_sr=target_sampling_rate
+        )
+        sampling_rate = target_sampling_rate
+
+    return {"array": list(array), "sampling_rate": sampling_rate}
+
+
+def _download_yt_video(url: str) -> str:
+    from pytube import YouTube
+
+    youtube = YouTube(url)
+    video = youtube.streams.first()
+
+    fn = "".join(random.choices("abcdefghijklmnopqrstuvwxyz", k=10))
+    file_path = video.download(output_path=tempfile.gettempdir(), filename=fn)
+
+    return file_path
+
+
+def _read_video_pyav(container, indices):
+    frames = []
+    container.seek(0)
+    start_index = indices[0]
+    end_index = indices[-1]
+    for i, frame in enumerate(container.decode(video=0)):
+        if i > end_index:
+            break
+        if i >= start_index and i in indices:
+            frames.append(frame)
+    return np.stack([x.to_ndarray(format="rgb24") for x in frames])
+
+
+def _sample_frame_indices(clip_len, frame_sample_rate, seg_len):
+    converted_len = int(clip_len * frame_sample_rate)
+    end_idx = np.random.randint(converted_len, seg_len)
+    start_idx = end_idx - converted_len
+    indices = np.linspace(start_idx, end_idx, num=clip_len)
+    indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
+    return indices
+
+
+def load_video(
+    input_: str,
+    frames: int = 8,
+    frame_sample_rate: int = 1,
+    start_time: Optional[int] = None,
+    end_time: Optional[int] = None,
+) -> np.ndarray:
+    import av
+
+    delete_file = False
+
+    if isinstance(input_, dict) and "youtube.com" and input_.get("url", ""):
+        file_path = _download_yt_video(input_["url"])
+        delete_file = True
+        # start_time = input_.get("start_time", None)
+        # end_time = input_.get("end_time", None)
+    elif isinstance(input_, str) and "youtube.com" in input_:
+        file_path = _download_yt_video(input_)
+        delete_file = True
+    elif isinstance(input_, str):
+        file_path = input_
+    else:
+        raise ValueError(f"Could not load video from {input_}")
+
+    if start_time is not None or end_time is not None:
+        start_time = start_time if start_time is not None else 0
+        end_time = end_time if end_time is not None else "end"
+        trim_file_path = f"{file_path.rsplit('.', 1)[0]}_trim.mp4"
+        subprocess.run(
+            [
+                "ffmpeg",
+                "-i",
+                file_path,
+                "-ss",
+                str(start_time),
+                "-to",
+                str(end_time),
+                "-c",
+                "copy",
+                trim_file_path,
+            ]
+        )
+        file_path = trim_file_path
+
+    container = av.open(file_path)
+    indices = _sample_frame_indices(
+        clip_len=frames,
+        frame_sample_rate=frame_sample_rate,
+        seg_len=container.streams.video[0].frames,
+    )
+    video = _read_video_pyav(container, indices)
+
+    if delete_file:
+        os.remove(file_path)
+
+    return video

src/sonicverse/multi_token/inference.py

@@ -0,0 +1,83 @@
+from typing import Type, List, Optional
+import logging
+
+from transformers import AutoTokenizer, AutoConfig, BitsAndBytesConfig
+from huggingface_hub import hf_hub_download
+from peft import PeftModel
+import torch
+import os
+
+from multi_token.model_utils import fix_tokenizer, MultiTaskType
+from multi_token.modalities.base_modality import Modality
+from multi_token.language_models.mistral import MistralForCausalLM
+from multi_token.language_models import LANGUAGE_MODEL_NAME_TO_CLASS
+from multi_token.modalities import MODALITY_BUILDERS
+
+
+def load_trained_lora_model(
+    model_name_or_path: str,
+    model_lora_path: str,
+    model_cls: Optional[Type] = None,
+    modalities: Optional[List[Modality]] = None,
+    load_bits: int = 16,
+    device_map: str = "auto",
+    use_multi_task: int = MultiTaskType.NO_MULTI_TASK,
+    tasks_config: str = None
+):
+    load_kwargs = {"device_map": device_map}
+
+    if load_bits == 8:
+        load_kwargs["load_in_8bit"] = True
+    elif load_bits == 4:
+        load_kwargs["load_in_4bit"] = True
+        load_kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+    elif load_bits == 16:
+        load_kwargs["torch_dtype"] = torch.float16
+    else:
+        raise ValueError(f"Invalid load_bits: {load_bits}")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
+    fix_tokenizer(tokenizer)
+
+    cfg = AutoConfig.from_pretrained(model_lora_path)
+    if model_cls is None:
+        model_cls = LANGUAGE_MODEL_NAME_TO_CLASS[cfg.model_cls]
+    if modalities is None:
+        if use_multi_task:
+            modalities = MODALITY_BUILDERS[cfg.modality_builder](use_multi_task = use_multi_task, tasks_config = tasks_config)
+        else:
+            modalities = MODALITY_BUILDERS[cfg.modality_builder]()
+
+    logging.info(f"Loading base model from {model_name_or_path} as {load_bits} bits")
+    model = model_cls.from_pretrained(
+        model_name_or_path, low_cpu_mem_usage=True, config=cfg, **load_kwargs
+    )
+    model.modalities = modalities
+
+    logging.info(f"Loading projector weights for {[m.name for m in modalities]}")
+    if os.path.exists(os.path.join(model_lora_path, "non_lora_trainables.bin")):
+        non_lora_trainables = torch.load(
+            os.path.join(model_lora_path, "non_lora_trainables.bin"), map_location="cuda"
+        )
+    else:
+        local_fn = hf_hub_download(
+            repo_id=model_lora_path,
+            filename="non_lora_trainables.bin",
+            repo_type="model",
+        )
+        non_lora_trainables = torch.load(local_fn, map_location="cuda")
+    model.get_model().initialize_pretrained_modules(modalities, non_lora_trainables)
+
+    logging.info(f"Loading and merging LoRA weights from {model_lora_path}")
+    model = PeftModel.from_pretrained(model, model_lora_path)
+    if load_bits == 16:
+        # TODO: Figure out why this fails for other bit sizes
+        model = model.merge_and_unload()
+    model.eval()
+
+    return model, tokenizer

src/sonicverse/multi_token/language_models/__init__.py

@@ -0,0 +1,7 @@
+from multi_token.language_models.mistral import (
+    MistralLMMForCausalLM,
+)
+
+LANGUAGE_MODEL_CLASSES = [MistralLMMForCausalLM]
+
+LANGUAGE_MODEL_NAME_TO_CLASS = {cls.__name__: cls for cls in LANGUAGE_MODEL_CLASSES}

src/sonicverse/multi_token/language_models/base_model.py

@@ -0,0 +1,181 @@
+from typing import List, Dict
+from abc import ABC, abstractmethod
+
+from torch.nn.functional import conv1d
+import torch
+import logging
+
+from multi_token.modalities.base_modality import Modality
+from multi_token.model_utils import MultiTaskType
+
+from torchviz import make_dot
+
+class LMMMetaModel:
+    def __init__(self, config):
+        super(LMMMetaModel, self).__init__(config)
+
+    def _load_projector_weights(self, weights: Dict):
+        weights = {
+            (k[23:] if k.startswith("base_model.model.model.") else k): v
+            for k, v in weights.items()
+        }
+        logging.info(f"Loading pretrained weights: {list(weights.keys())}")
+        load_result = self.load_state_dict(weights, strict=False)
+        assert (
+            len(load_result.unexpected_keys) == 0
+        ), "Unexpected weights, is this the right model?"
+
+    def initialize_pretrained_modules(self, modalities: List[Modality], weights: Dict):
+        for m in modalities:
+            # projector = m.build_projector(self.config.hidden_size)
+            # setattr(self, m.name + "_lmm_projector", projector)
+            projector = m.build_projector(self.config.hidden_size)
+            if m.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+                for task_name in m.tasks["task_heads"].keys():
+                    task_model = projector[task_name]
+                    setattr(self, m.name + "_" + task_name, task_model)
+            else:
+                setattr(self, m.name + "_lmm_projector", projector)
+
+        self._load_projector_weights(weights)
+
+    def initialize_modules(self, modalities: List[Modality], weights: Dict):
+        names = [m.name for m in modalities]
+
+        self.config.modalities = names
+
+        for m in modalities:
+            # projector = m.build_projector(self.config.hidden_size)
+            # setattr(self, m.name + "_lmm_projector", projector)
+            projector = m.build_projector(self.config.hidden_size)
+            if m.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+                for task_name in m.tasks["task_heads"].keys():
+                    task_model = projector[task_name]
+                    setattr(self, m.name + "_" + task_name, task_model)
+            else:
+                setattr(self, m.name + "_lmm_projector", projector)
+
+        self._load_projector_weights(weights)
+
+
+class LMMMetaForCausalLM(ABC):
+    @abstractmethod
+    def get_model(self) -> "LMMMetaForCausalLM":
+        pass
+
+    def prepare_inputs_labels_for_multimodal(
+        self, input_ids, attention_mask, past_key_values, labels, **kwargs
+    ):
+        model = self.get_model()
+
+        batch_size, seq_len = input_ids.shape
+
+        # batch_size x seq_len x embedding_hidden_size
+        inputs_embeds = torch.zeros(
+            (batch_size, seq_len, self.config.hidden_size),
+            dtype=self.dtype,
+            device=self.device,
+        )
+
+        # modality x batch_size x instance_idx x modality_token_width x embedding_hidden_size
+        projected_tensors = []
+        # assuming that if caching is enabled, we'll never have past_key_values AND need to encode the instruction modality values
+        task_vals = {}
+
+        #print("here past_key_values", past_key_values)
+        #past_key_values == None
+        if past_key_values is None:
+            for m in self.modalities:
+                m_vals = m.forward(kwargs.get(m.name))
+                mp_vals = []
+                if m.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+                    proj = {}
+                    for task_name in m.tasks["task_heads"].keys():
+                        proj[task_name] = getattr(model, m.name + "_" + task_name)
+                else:
+                    proj = getattr(model, m.name + "_lmm_projector")
+
+                # project each batch into language model token space
+                for m_val in m_vals:
+                    if m.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+                        for task_name in m.tasks["task_heads"].keys():
+                            if task_name == "lmm_projector":
+                                mp_vals.append(proj[task_name](m_val))
+                                # make_dot(mp_vals[-1], params=dict(list(model.named_parameters()))).render(task_name, format="png")
+                            else:
+                                if task_name not in task_vals:
+                                    task_vals[task_name] = [proj[task_name](m_val)]
+                                else:
+                                    task_vals[task_name].append(proj[task_name](m_val))
+                                # make_dot(task_vals[task_name], params=dict(list(model.named_parameters()))).render(task_name, format="png")
+
+                    elif m.use_multi_task == MultiTaskType.PROJECTED_MULTI_TASK:
+                        task_outputs = proj(m_val)
+                        mp_vals.append(task_outputs.pop("projectors"))
+                        for task_name in task_outputs.keys():
+                            if not task_name in task_vals:
+                                task_vals[task_name] = [task_outputs[task_name]]
+                            else:
+                                task_vals[task_name].append(task_outputs[task_name])
+                    else:
+                        mp_vals.append(proj(m_val))
+
+                assert all(
+                    mp_val.shape[1:] == (m.token_width, self.config.hidden_size)
+                    for mp_val in mp_vals
+                ), (
+                    "Modality tensors have incorrect shape, check your projector implementation "
+                    + str([mp_val.shape[1:] for mp_val in mp_vals])
+                    + " vs expected "
+                    + str((m.token_width, self.config.hidden_size))
+                )
+                projected_tensors.append(mp_vals)
+
+        indices = None
+        for i, input_ids_sample in enumerate(input_ids):
+            is_text_mask = input_ids_sample >= 0
+
+            # fill in all the LLM-based text embeddings
+            inputs_embeds[i, is_text_mask] = model.embed_tokens(
+                input_ids_sample[is_text_mask]
+            )
+
+            # skip if all tokens are text tokens
+            if is_text_mask.sum() == seq_len:
+                continue
+            assert (
+                past_key_values is None
+            ), "We shouldn't have cached keys if this is the first instruction pass"
+
+            #past_key_values = None
+
+            for mi, m in enumerate(self.modalities):
+                # locate the group of tokens for this modality
+                m_mask = (input_ids_sample == m.token_idx).float()
+                m_kernel = torch.tensor(
+                    [-1] * m.token_width, dtype=m_mask.dtype, device=m_mask.device
+                )
+                m_conv = conv1d(
+                    m_mask.unsqueeze(0).unsqueeze(0),
+                    m_kernel.unsqueeze(0).unsqueeze(0),
+                )
+
+                # where do we see `token_width`-tokens in a row?
+                indices = (m_conv[0, 0] == -m.token_width).nonzero(as_tuple=True)[0]
+
+                # fill these embeddings with the projected modality tensor
+                last_covered_idx = -1
+                k = 0
+                for possible_token_idx in indices:
+                    if possible_token_idx <= last_covered_idx:
+                        # make sure we don't overwrite an instance we've already covered
+                        # handles bug caused by back-to-back tokens
+                        continue
+                    batch_modality_tensor = projected_tensors[mi][i][k]
+                    inputs_embeds[
+                        i, possible_token_idx : possible_token_idx + m.token_width
+                    ] = batch_modality_tensor
+                    last_covered_idx = possible_token_idx + m.token_width - 1
+                    k += 1
+
+        return None, attention_mask, past_key_values, inputs_embeds, labels, task_vals

src/sonicverse/multi_token/language_models/mistral.py

@@ -0,0 +1,235 @@
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from transformers import (
+    AutoConfig,
+    AutoModelForCausalLM,
+    MistralConfig,
+    MistralModel,
+    MistralForCausalLM,
+)
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from multi_token.language_models.base_model import (
+    LMMMetaModel,
+    LMMMetaForCausalLM,
+)
+
+
+class MistralLMMConfig(MistralConfig):
+    model_type = "mistral-lmm"
+
+
+class MistralLMMModel(LMMMetaModel, MistralModel):
+    config_class = MistralLMMConfig
+
+    def __init__(self, config: MistralLMMConfig):
+        super(MistralLMMModel, self).__init__(config)
+
+
+class MistralLMMForCausalLM(MistralForCausalLM, LMMMetaForCausalLM):
+    config_class = MistralLMMConfig
+
+    def __init__(self, config):
+        super(MistralForCausalLM, self).__init__(config)
+        self.model = MistralLMMModel(config)
+
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.modalities = None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self) -> "MistralLMMForCausalLM":
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[List] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        #print("Past keys ",past_key_values)
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if labels != None:
+            labels_inp = labels[0]
+        else:
+            labels_inp = labels
+        (
+            input_ids,
+            attention_mask,
+            past_key_values,
+            inputs_embeds,
+            lmm_labels,
+            task_values
+        ) = self.prepare_inputs_labels_for_multimodal(
+            input_ids, attention_mask, past_key_values, labels_inp, **kwargs
+        )
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        # print("Labels 1 size ", len(labels[1]))
+        # print("labels 1 element size ", len(labels[1][0]))
+        # print("labels 1 element 1 task size ", labels[1][0][0].shape)
+        # print("labels 1 element 2 task size ", labels[1][0][1].shape)
+        # print("labels 1 element 3 task size ", labels[1][0][2].shape)
+        # print("task vals size ", len(task_values))
+        # for task in task_values.keys():
+        #     print(" task ", task, len(task_values[task]))
+        #     print(" task element", task, task_values[task][0].shape)
+
+
+        if labels != None:
+            task_pairs = {}
+            task_list = list(task_values.keys())
+            for task_id in range(len(task_list)):
+                _task_labels = []
+                _task_outputs = []
+
+                _task = task_list[task_id]
+                for inst in range(len(task_values[_task])):
+                    # print("task output shape ", _task, task_values[_task][inst].shape)
+                    _task_outputs.append(task_values[_task][inst].unsqueeze(0))
+                    _task_labels.append(torch.stack([labels[1][inst][task_id]]))
+
+                task_pairs[_task] = [_task_labels, _task_outputs]
+                # print("TASK ", _task)
+                # print(" LABELS LEN ", len(task_pairs[_task][0]))
+                # print(" LABELS ELEM shape ", task_pairs[_task][0][0].shape)
+                # print(" VALUES LEN ", len(task_pairs[_task][1]))
+                # print(" VALUES ELEM shape ", task_pairs[_task][1][0].shape)
+
+        loss = None
+        if lmm_labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = lmm_labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        # print("loss ", loss)
+
+
+        if labels != None:
+            task_loss = {}
+            for task in task_list:
+                preds = torch.cat(task_pairs[task][1], dim=0)
+                labs = torch.cat(task_pairs[task][0], dim=0)
+                preds_flat = preds.view(-1, preds.size(-1))  # Reshape to (batch_size * sequence_length, num_classes)
+                labs_flat = labs.view(-1)  # Reshape to (batch_size * sequence_length)
+
+                #print("task ", task)
+                #print("preds shape ", preds.shape)
+                #print("labs shape ", labs.shape)
+                if task == "lmm_projector":
+                    task_loss[task] = CrossEntropyLoss()(preds,labs)
+                else:
+                    task_loss[task] = nn.BCEWithLogitsLoss()(preds, labs)
+        # print("task losses ", task_loss)
+
+        total_loss = None
+        if labels != None:
+            total_task_loss = None
+            for task in task_list:
+                if self.modalities[0].tasks["task_heads"][task]["weight"] != 0.0:
+                    if total_task_loss != None:
+                        total_task_loss += self.modalities[0].tasks["task_heads"][task]["weight"]*task_loss[task]
+                    else:
+                        total_task_loss = self.modalities[0].tasks["task_heads"][task]["weight"]*task_loss[task]
+
+            if total_task_loss != None:
+                total_loss = self.modalities[0].tasks["task_heads"]["lmm_projector"]["weight"]*loss + total_task_loss
+            else:
+                total_loss = loss
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (total_loss,) + output if total_loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=total_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        modality_inputs=None,
+        **kwargs
+    ):
+        #print("hoooo", past_key_values)
+
+        #past_key_values = None
+        if past_key_values:
+            input_ids = input_ids[:, -1:]
+
+        if inputs_embeds is not None:
+            raise ValueError("inputs_embeds not supported")
+
+        model_inputs = {
+            "input_ids": input_ids,
+            "position_ids": None,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "attention_mask": attention_mask,
+            **(modality_inputs or {}),
+        }
+
+        return model_inputs
+
+
+AutoConfig.register("mistral-lmm", MistralLMMConfig)
+AutoModelForCausalLM.register(MistralLMMConfig, MistralLMMForCausalLM)

src/sonicverse/multi_token/modalities/__init__.py

@@ -0,0 +1,31 @@
+from multi_token.model_utils import MultiTaskType
+from multi_token.modalities.vision_clip import (
+    CLIPVisionModality,
+    OUTPUT_LAYER as CLIP_POOL_LAYER,
+)
+from multi_token.modalities.imagebind import ImageBindModality
+from multi_token.modalities.document_gte import DocumentGTEModality
+from multi_token.modalities.audio_whisper import WhisperAudioModality
+from multi_token.modalities.audio_clap import CLAPAudioModality
+from multi_token.modalities.video_xclip import XCLIPVideoModality
+from multi_token.modalities.audio_descript import DescriptAudioModality
+from multi_token.modalities.audio_mert import MERTAudioModality
+
+MODALITY_BUILDERS = {
+    "vision_clip": lambda: [CLIPVisionModality()],
+    "vision_clip_pool": lambda: [
+        CLIPVisionModality(feature_layer=CLIP_POOL_LAYER, num_tokens_output=10)
+    ],
+    "audio_whisper": lambda: [
+        WhisperAudioModality(
+            num_tokens_output=10, model_name_or_path="openai/whisper-small"
+        )
+    ],
+    "audio_mert": lambda use_multi_task=MultiTaskType.NO_MULTI_TASK, tasks_config=None :[MERTAudioModality(use_multi_task=use_multi_task, tasks_config=tasks_config, num_tokens_output=60, hidden_dim=32, num_conv_layers = 3, num_mlp_layers = 2)],
+    "audio_clap": lambda use_multi_task=MultiTaskType.NO_MULTI_TASK, tasks_config=None :[CLAPAudioModality(use_multi_task=use_multi_task, tasks_config=tasks_config, num_tokens_output=20)],
+    "audio_descript": lambda use_multi_task=MultiTaskType.NO_MULTI_TASK, tasks_config=None : [DescriptAudioModality(use_multi_task=use_multi_task, tasks_config=tasks_config, num_projector_conv_layers=1, num_projector_mlp_layers=1, num_tokens_output=60, codebooks=96)],
+    "video_xclip": lambda: [XCLIPVideoModality(num_tokens_output=10)],
+    "imagebind": lambda: [ImageBindModality()],
+    "document_gte": lambda: [DocumentGTEModality()],
+    "document_gte_x16": lambda: [DocumentGTEModality(num_tokens_output=32)],
+}

src/sonicverse/multi_token/modalities/audio_clap.py

@@ -0,0 +1,142 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from transformers import ClapModel, ClapProcessor
+
+from multi_token.model_utils import MultiTaskType
+from multi_token.data_tools import load_audio
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_mlp_vector_projector, build_mt_vector_projector, MultiTaskModel
+)
+
+import json
+
+OUTPUT_EMB_SIZE = 512
+
+
+class CLAPAudioModule(nn.Module):
+    def __init__(self, model_name_or_path: str):
+        super().__init__()
+        self.model_name_or_path = model_name_or_path
+        self.model = None
+        self.processor = None
+
+        self.load_model()
+
+    def load_model(self):
+        self.model = ClapModel.from_pretrained(self.model_name_or_path)
+        self.processor = ClapProcessor.from_pretrained(self.model_name_or_path)
+        self.model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, audios) -> torch.Tensor:
+        embs = []
+        for audio_features in audios:
+            features = self.model.get_audio_features(
+                input_features=audio_features["input_features"].to(torch.float32),
+                is_longer=audio_features["is_longer"],
+            )
+            embs.append(features)
+        embs = torch.stack(embs)
+        return embs.view(-1, 1, OUTPUT_EMB_SIZE)
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def device(self):
+        return self.model.device
+
+
+class CLAPAudioModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = "laion/clap-htsat-fused",
+        num_projector_layers: int = 2,
+        num_tokens_output: int = 10,
+        use_multi_task: int = MultiTaskType.NO_MULTI_TASK,
+        tasks_config: str = None
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = CLAPAudioModule(model_name_or_path=self.model_name_or_path)
+        self.num_projector_layers = num_projector_layers
+        self.num_tokens_output = num_tokens_output
+        self.dtype = torch.float32
+        self.use_multi_task = use_multi_task
+        self.tasks = None
+        if self.use_multi_task != MultiTaskType.NO_MULTI_TASK:            
+            with open(tasks_config, 'r') as f:
+                self.tasks = json.load(f)
+
+        print("Tasks :", self.tasks)
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        if self.use_multi_task == MultiTaskType.PROJECTED_MULTI_TASK:
+            return MultiTaskModel(OUTPUT_EMB_SIZE, self.tasks)
+        elif self.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+            return build_mt_vector_projector(
+            # return build_mlp_vector_projector(
+                input_hidden_size=OUTPUT_EMB_SIZE,
+                lm_hidden_size=lm_hidden_size,
+            #     num_layers=self.num_projector_layers,
+            #     num_tokens=self.num_tokens_output,
+            # )
+                tasks = self.tasks
+            )
+            # )["llm_projector"]
+        else:
+            return build_mlp_vector_projector(
+                input_hidden_size=OUTPUT_EMB_SIZE,
+                lm_hidden_size=lm_hidden_size,
+                num_layers=self.num_projector_layers,
+                num_tokens=self.num_tokens_output,
+            )
+
+    @property
+    def name(self) -> str:
+        return "audio_clap"
+
+    @property
+    def token(self) -> str:
+        return "<sound>"
+
+    @property
+    def data_key(self) -> str:
+        return "sounds"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "CLAPAudioModality":
+        self.dtype = dtype
+        self.module.to(device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[Dict]]:
+        row_values = []
+        for row in rows:
+            audios = []
+            for audio_dict in row[self.data_key]:
+                audio_dict = load_audio(
+                    audio_dict,
+                    target_sampling_rate=self.module.processor.feature_extractor.sampling_rate,
+                )
+                audio_processed = self.module.processor(
+                    audios=audio_dict["array"],
+                    return_tensors="pt",
+                    sampling_rate=audio_dict["sampling_rate"],
+                )
+                audios.append(audio_processed)
+            row_values.append(audios)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        audio_features = []
+        for audio_batch in encoded_values:
+            audio_features.append(self.module.forward(audio_batch).to(dtype=self.dtype))
+        return audio_features

src/sonicverse/multi_token/modalities/audio_descript.py

@@ -0,0 +1,169 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import dac
+from audiotools import AudioSignal
+
+from multi_token.model_utils import MultiTaskType
+from multi_token.data_tools import load_audio_signal
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_mlp_vector_projector, build_attentive_cnn_projector, build_cnn_mlp_projector, MultiTaskModel
+)
+
+import json
+
+OUTPUT_FRAMES_SIZE = 512
+# OUTPUT_EMB_SIZE = 2048
+OUTPUT_EMB_CHANNELS = 96
+
+class DescriptAudioModule(nn.Module):
+    def __init__(self, model_name_or_path: str, codebooks = 4):
+        super().__init__()
+        self.model_name_or_path = model_name_or_path
+        self.model = None
+        self.processor = None
+        self.codebooks = codebooks
+
+        self.load_model()
+
+    def load_model(self):
+        # self.model = ClapModel.from_pretrained(self.model_name_or_path)
+        self.model = dac.DAC.load(self.model_name_or_path)
+
+    def forward(self, audios) -> torch.Tensor:
+        embs = []
+        for audio_features in audios:
+            # print("Audio features sample rate ", audio_features[0].sample_rate)
+            x = self.model.preprocess(audio_features[0].audio_data, audio_features[0].sample_rate)
+            z, codes, latents, _, _ = self.model.encode(x)
+
+            # print("latents og shape ", latents.shape)
+            # If the tensor is larger than desired_shape, crop it
+            if latents.shape[2] > OUTPUT_FRAMES_SIZE:
+                latents = latents[:, :, :OUTPUT_FRAMES_SIZE]
+            # If the tensor is smaller than desired_shape, pad it
+            elif latents.shape[2] < OUTPUT_FRAMES_SIZE:
+                pad_width = (0, OUTPUT_FRAMES_SIZE - latents.shape[2])
+                latents = torch.nn.functional.pad(latents, pad_width)
+                # print("Codes new shape ", codes_new.shape)
+
+            # print("latents int shape ", latents.shape)
+
+            latents = latents[0][:self.codebooks]
+            
+            # print("latents final shape ", latents.shape)
+
+            embs.append(latents)
+
+        embs = torch.stack(embs)
+
+        # output_embs = embs.view(-1, 1, OUTPUT_FRAMES_SIZE*self.codebooks)
+        # print("embs post view shape ", output_embs.shape)
+
+        return embs
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def device(self):
+        return self.model.device
+
+
+class DescriptAudioModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = dac.utils.download(model_type="16khz"),
+        num_projector_conv_layers: int = 2,
+        num_projector_mlp_layers: int = 2,
+        num_tokens_output: int = 10,
+        codebooks: int = 96,
+        use_multi_task: MultiTaskType = MultiTaskType.NO_MULTI_TASK,
+        tasks_config: str = None
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = DescriptAudioModule(model_name_or_path=self.model_name_or_path, codebooks=codebooks)
+        self.num_projector_conv_layers = num_projector_conv_layers
+        self.num_projector_mlp_layers = num_projector_mlp_layers
+        self.num_tokens_output = num_tokens_output
+        self.dtype = torch.float32
+        self.codebooks = codebooks
+        self.use_multi_task = use_multi_task
+        self.tasks = None
+        if self.use_multi_task != MultiTaskType.NO_MULTI_TASK:            
+            with open(tasks_config, 'r') as f:
+                self.tasks = json.load(f)
+
+        print("Tasks :", self.tasks)
+
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        if self.use_multi_task == MultiTaskType.PROJECTED_MULTI_TASK:
+            projector = MultiTaskModel(OUTPUT_EMB_CHANNELS, 1, True, -1, False, self.tasks)
+            print("projector ", projector)
+            return projector
+        elif self.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+            return build_mt_vector_projector(
+            # return build_mlp_vector_projector(
+                input_hidden_size=OUTPUT_EMB_CHANNELS,
+                lm_hidden_size=lm_hidden_size,
+            #     num_layers=self.num_projector_layers,
+            #     num_tokens=self.num_tokens_output,
+            # )
+                tasks = self.tasks
+            )
+            # )["llm_projector"]
+        else:
+            return build_multi_layer_cnn_mlp_projector(
+                input_channels = OUTPUT_EMB_CHANNELS,
+                input_size = OUTPUT_EMB_SIZE,
+                num_feature_layers= OUTPUT_FEATURE_LAYERS,
+                lm_hidden_size = lm_hidden_size,
+                num_tokens = self.num_tokens_output,
+                hidden_dim = self.hidden_dim,
+                num_conv_layers = self.num_conv_layers,
+                num_mlp_layers = self.num_mlp_layers
+            )
+
+    @property
+    def name(self) -> str:
+        return "audio_descript"
+
+    @property
+    def token(self) -> str:
+        return "<sound>"
+
+    @property
+    def data_key(self) -> str:
+        return "sounds"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "DescriptAudioModality":
+        self.dtype = dtype
+        self.module.to(device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[Dict]]:
+        row_values = []
+        for row in rows:
+            audios = []
+            for audio_dict in row[self.data_key]:
+                audio_dict = load_audio_signal(
+                    audio_dict
+                )
+                audios.append(audio_dict["array"])
+            row_values.append(audios)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        audio_features = []
+        for audio_batch in encoded_values:
+            audio_features.append(self.module.forward(audio_batch).to(dtype=self.dtype))
+        return audio_features

src/sonicverse/multi_token/modalities/audio_descript_bu.py

@@ -0,0 +1,133 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import dac
+from audiotools import AudioSignal
+
+
+from multi_token.data_tools import load_audio_signal
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_mlp_vector_projector, build_attentive_cnn_projector, build_cnn_mlp_projector
+)
+
+OUTPUT_FRAMES_SIZE = 512
+# OUTPUT_EMB_SIZE = 2048
+
+class DescriptAudioModule(nn.Module):
+    def __init__(self, model_name_or_path: str, codebooks = 4):
+        super().__init__()
+        self.model_name_or_path = model_name_or_path
+        self.model = None
+        self.processor = None
+        self.codebooks = codebooks
+
+        self.load_model()
+
+    def load_model(self):
+        # self.model = ClapModel.from_pretrained(self.model_name_or_path)
+        self.model = dac.DAC.load(self.model_name_or_path)
+
+    def forward(self, audios) -> torch.Tensor:
+        embs = []
+        for audio_features in audios:
+            x = self.model.preprocess(audio_features[0].audio_data, audio_features[0].sample_rate)
+            z, codes, latents, _, _ = self.model.encode(x)
+
+            # If the tensor is larger than desired_shape, crop it
+            if codes.shape[2] > OUTPUT_FRAMES_SIZE:
+                codes = codes[:, :, :OUTPUT_FRAMES_SIZE]
+            # If the tensor is smaller than desired_shape, pad it
+            elif codes.shape[2] < OUTPUT_FRAMES_SIZE:
+                pad_width = (0, OUTPUT_FRAMES_SIZE - codes.shape[2])
+                codes = torch.nn.functional.pad(codes, pad_width)
+                # print("Codes new shape ", codes_new.shape)
+
+            codes_of_interest = codes[0][:self.codebooks]
+            
+            embs.append(codes_of_interest)
+
+        embs = torch.stack(embs)
+
+        # output_embs = embs.view(-1, 1, OUTPUT_FRAMES_SIZE*self.codebooks)
+        # print("embs post view shape ", output_embs.shape)
+
+        return embs
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def device(self):
+        return self.model.device
+
+
+class DescriptAudioModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = dac.utils.download(model_type="16khz"),
+        num_projector_conv_layers: int = 2,
+        num_projector_mlp_layers: int = 2,
+        num_tokens_output: int = 10,
+        codebooks: int = 4
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = DescriptAudioModule(model_name_or_path=self.model_name_or_path, codebooks=codebooks)
+        self.num_projector_conv_layers = num_projector_conv_layers
+        self.num_projector_mlp_layers = num_projector_mlp_layers
+        self.num_tokens_output = num_tokens_output
+        self.dtype = torch.float32
+        self.codebooks = codebooks
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        return build_cnn_mlp_projector(
+            input_channels=self.codebooks,
+            input_size=OUTPUT_FRAMES_SIZE,
+            lm_hidden_size=lm_hidden_size,
+            num_tokens=self.num_tokens_output,
+            hidden_dim=64,
+            num_conv_layers=self.num_projector_conv_layers,
+            num_mlp_layers=self.num_projector_mlp_layers
+        )
+
+    @property
+    def name(self) -> str:
+        return "audio_descript"
+
+    @property
+    def token(self) -> str:
+        return "<sound>"
+
+    @property
+    def data_key(self) -> str:
+        return "sounds"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "DescriptAudioModality":
+        self.dtype = dtype
+        self.module.to(device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[Dict]]:
+        row_values = []
+        for row in rows:
+            audios = []
+            for audio_dict in row[self.data_key]:
+                audio_dict = load_audio_signal(
+                    audio_dict
+                )
+                audios.append(audio_dict["array"])
+            row_values.append(audios)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        audio_features = []
+        for audio_batch in encoded_values:
+            audio_features.append(self.module.forward(audio_batch).to(dtype=self.dtype))
+        return audio_features

src/sonicverse/multi_token/modalities/audio_mert.py

@@ -0,0 +1,162 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from transformers import Wav2Vec2FeatureExtractor, AutoModel
+
+from multi_token.model_utils import MultiTaskType
+from multi_token.data_tools import load_audio
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_mlp_vector_projector, build_mt_vector_projector, build_multi_layer_cnn_mlp_projector, MultiTaskModel
+)
+from multi_token.modalities.multi_task_projector_shared import MultiTaskSharedModel
+
+import json
+
+OUTPUT_EMB_CHANNELS = 768 #1024
+OUTPUT_EMB_SIZE = 760
+OUTPUT_FEATURE_LAYERS = 13 #25
+
+cache_dir="/home/ubuntu/.cache/"
+
+class MERTAudioModule(nn.Module):
+    def __init__(self, model_name_or_path: str):
+        super().__init__()
+        self.model_name_or_path = model_name_or_path
+        self.model = None
+        self.processor = None
+
+        self.load_model()
+
+    def load_model(self):
+        self.model = AutoModel.from_pretrained(self.model_name_or_path, trust_remote_code=True, cache_dir=cache_dir)
+        self.processor = Wav2Vec2FeatureExtractor.from_pretrained(self.model_name_or_path,trust_remote_code=True, cache_dir=cache_dir)
+        self.model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, audios) -> torch.Tensor:
+        embs = []
+        for audio_features in audios:
+            outputs = self.model(**audio_features.to(torch.float32), output_hidden_states=True)
+            features = torch.stack(outputs.hidden_states).squeeze()
+            embs.append(features)
+        embs = torch.stack(embs)
+        embs = embs.squeeze()
+        padding_needed = OUTPUT_EMB_SIZE - embs.shape[1]
+        embs = torch.nn.functional.pad(embs, (0, 0, 0, padding_needed, 0, 0))
+        return embs
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def device(self):
+        return self.model.device
+
+
+class MERTAudioModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = "m-a-p/MERT-v1-95M",
+        num_tokens_output: int = 10,
+        hidden_dim: int = 32,
+        num_conv_layers: int = 5,
+        num_mlp_layers: int = 5,
+        use_multi_task: MultiTaskType = MultiTaskType.NO_MULTI_TASK,
+        tasks_config: str = None
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = MERTAudioModule(model_name_or_path=self.model_name_or_path)
+        self.num_tokens_output = num_tokens_output
+        self.hidden_dim = hidden_dim
+        self.num_conv_layers = num_conv_layers
+        self.num_mlp_layers = num_mlp_layers
+        self.dtype = torch.float32
+        self.use_multi_task = use_multi_task
+        self.tasks = None
+        if self.use_multi_task != MultiTaskType.NO_MULTI_TASK:            
+            with open(tasks_config, 'r') as f:
+                self.tasks = json.load(f)
+
+        print("Tasks :", self.tasks)
+
+    # all_layer_hidden_states = torch.stack(outputs.hidden_states).squeeze()
+    # print(all_layer_hidden_states.shape) # [25 layer, Time steps, 1024 feature_dim]
+    # time_reduced_hidden_states = all_layer_hidden_states.mean(-2)
+    # print(time_reduced_hidden_states.shape) # [25, 1024]
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        if self.use_multi_task == MultiTaskType.PROJECTED_MULTI_TASK:
+            projector = MultiTaskSharedModel(self.tasks)
+            print("projector ", projector)
+            return projector
+        elif self.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+            return build_mt_vector_projector(
+            # return build_mlp_vector_projector(
+                input_hidden_size=OUTPUT_EMB_SIZE,
+                lm_hidden_size=lm_hidden_size,
+            #     num_layers=self.num_projector_layers,
+            #     num_tokens=self.num_tokens_output,
+            # )
+                tasks = self.tasks
+            )
+            # )["llm_projector"]
+        else:
+            return build_multi_layer_cnn_mlp_projector(
+                input_channels = OUTPUT_EMB_CHANNELS,
+                input_size = OUTPUT_EMB_SIZE,
+                num_feature_layers= OUTPUT_FEATURE_LAYERS,
+                lm_hidden_size = lm_hidden_size,
+                num_tokens = self.num_tokens_output,
+                hidden_dim = self.hidden_dim,
+                num_conv_layers = self.num_conv_layers,
+                num_mlp_layers = self.num_mlp_layers
+            )
+
+    @property
+    def name(self) -> str:
+        return "audio_mert"
+
+    @property
+    def token(self) -> str:
+        return "<sound>"
+
+    @property
+    def data_key(self) -> str:
+        return "sounds"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "MERTAudioModality":
+        self.dtype = dtype
+        self.module.to(device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[Dict]]:
+        row_values = []
+        for row in rows:
+            audios = []
+            for audio_dict in row[self.data_key]:
+                audio_dict = load_audio(
+                    audio_dict,
+                    target_sampling_rate=self.module.processor.sampling_rate,
+                )
+                audio_processed = self.module.processor(
+                    audio_dict["array"],
+                    return_tensors="pt",
+                    sampling_rate=audio_dict["sampling_rate"],
+                )
+                audios.append(audio_processed)
+            row_values.append(audios)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        audio_features = []
+        for audio_batch in encoded_values:
+            audio_features.append(self.module.forward(audio_batch).to(dtype=self.dtype))
+        return audio_features

src/sonicverse/multi_token/modalities/audio_mert_bu.py

@@ -0,0 +1,159 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from transformers import Wav2Vec2FeatureExtractor, AutoModel
+
+from multi_token.model_utils import MultiTaskType
+from multi_token.data_tools import load_audio
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_mlp_vector_projector, build_mt_vector_projector, build_multi_layer_cnn_mlp_projector, MultiTaskModel
+)
+
+import json
+
+OUTPUT_EMB_CHANNELS = 1024
+OUTPUT_EMB_SIZE = 760
+OUTPUT_FEATURE_LAYERS = 25
+
+class MERTAudioModule(nn.Module):
+    def __init__(self, model_name_or_path: str):
+        super().__init__()
+        self.model_name_or_path = model_name_or_path
+        self.model = None
+        self.processor = None
+
+        self.load_model()
+
+    def load_model(self):
+        self.model = AutoModel.from_pretrained(self.model_name_or_path, trust_remote_code=True)
+        self.processor = Wav2Vec2FeatureExtractor.from_pretrained(self.model_name_or_path,trust_remote_code=True)
+        self.model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, audios) -> torch.Tensor:
+        embs = []
+        for audio_features in audios:
+            outputs = self.model(**audio_features.to(torch.float32), output_hidden_states=True)
+            features = torch.stack(outputs.hidden_states).squeeze()
+            embs.append(features)
+        embs = torch.stack(embs)
+        embs = embs.squeeze()
+        padding_needed = OUTPUT_EMB_SIZE - embs.shape[1]
+        embs = torch.nn.functional.pad(embs, (0, 0, 0, padding_needed, 0, 0))
+        return embs
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def device(self):
+        return self.model.device
+
+
+class MERTAudioModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = "m-a-p/MERT-v1-330M",
+        num_tokens_output: int = 10,
+        hidden_dim: int = 32,
+        num_conv_layers: int = 5,
+        num_mlp_layers: int = 5,
+        use_multi_task: MultiTaskType = MultiTaskType.NO_MULTI_TASK,
+        tasks_config: str = None
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = MERTAudioModule(model_name_or_path=self.model_name_or_path)
+        self.num_tokens_output = num_tokens_output
+        self.hidden_dim = hidden_dim
+        self.num_conv_layers = num_conv_layers
+        self.num_mlp_layers = num_mlp_layers
+        self.dtype = torch.float32
+        self.use_multi_task = use_multi_task
+        self.tasks = None
+        if self.use_multi_task != MultiTaskType.NO_MULTI_TASK:            
+            with open(tasks_config, 'r') as f:
+                self.tasks = json.load(f)
+
+        print("Tasks :", self.tasks)
+
+    # all_layer_hidden_states = torch.stack(outputs.hidden_states).squeeze()
+    # print(all_layer_hidden_states.shape) # [25 layer, Time steps, 1024 feature_dim]
+    # time_reduced_hidden_states = all_layer_hidden_states.mean(-2)
+    # print(time_reduced_hidden_states.shape) # [25, 1024]
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        if self.use_multi_task == MultiTaskType.PROJECTED_MULTI_TASK:
+            projector = MultiTaskModel(OUTPUT_EMB_CHANNELS, OUTPUT_FEATURE_LAYERS, True, self.tasks)
+            print("projector ", projector)
+            return projector
+        elif self.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+            return build_mt_vector_projector(
+            # return build_mlp_vector_projector(
+                input_hidden_size=OUTPUT_EMB_SIZE,
+                lm_hidden_size=lm_hidden_size,
+            #     num_layers=self.num_projector_layers,
+            #     num_tokens=self.num_tokens_output,
+            # )
+                tasks = self.tasks
+            )
+            # )["llm_projector"]
+        else:
+            return build_multi_layer_cnn_mlp_projector(
+                input_channels = OUTPUT_EMB_CHANNELS,
+                input_size = OUTPUT_EMB_SIZE,
+                num_feature_layers= OUTPUT_FEATURE_LAYERS,
+                lm_hidden_size = lm_hidden_size,
+                num_tokens = self.num_tokens_output,
+                hidden_dim = self.hidden_dim,
+                num_conv_layers = self.num_conv_layers,
+                num_mlp_layers = self.num_mlp_layers
+            )
+
+    @property
+    def name(self) -> str:
+        return "audio_mert"
+
+    @property
+    def token(self) -> str:
+        return "<sound>"
+
+    @property
+    def data_key(self) -> str:
+        return "sounds"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "MERTAudioModality":
+        self.dtype = dtype
+        self.module.to(device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[Dict]]:
+        row_values = []
+        for row in rows:
+            audios = []
+            for audio_dict in row[self.data_key]:
+                audio_dict = load_audio(
+                    audio_dict,
+                    target_sampling_rate=self.module.processor.sampling_rate,
+                )
+                audio_processed = self.module.processor(
+                    audio_dict["array"],
+                    return_tensors="pt",
+                    sampling_rate=audio_dict["sampling_rate"],
+                )
+                audios.append(audio_processed)
+            row_values.append(audios)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        audio_features = []
+        for audio_batch in encoded_values:
+            audio_features.append(self.module.forward(audio_batch).to(dtype=self.dtype))
+        return audio_features

src/sonicverse/multi_token/modalities/audio_whisper.py

@@ -0,0 +1,120 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from transformers import AutoFeatureExtractor, WhisperModel
+
+from multi_token.data_tools import load_audio
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_mlp_vector_projector,
+)
+
+
+OUTPUT_EMB_SIZE = 768
+
+
+class WhisperAudioModule(nn.Module):
+    def __init__(self, model_name_or_path: str):
+        super().__init__()
+        self.model_name_or_path = model_name_or_path
+        self.model = None
+        self.feature_extractor = None
+
+        self.load_model()
+
+    def load_model(self):
+        self.model = WhisperModel.from_pretrained(self.model_name_or_path)
+        self.feature_extractor = AutoFeatureExtractor.from_pretrained(
+            self.model_name_or_path
+        )
+        self.model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, audios) -> torch.Tensor:
+        hidden_states = []
+        for i in range(audios.shape[0]):
+            decoder_input_ids = (
+                torch.tensor([[1]]) * self.model.config.decoder_start_token_id
+            )
+            last_hidden_state = self.model(
+                audios[i].to(device=self.device, dtype=self.dtype),
+                decoder_input_ids=decoder_input_ids.to(device=self.device),
+            ).last_hidden_state
+            hidden_states.append(last_hidden_state)
+        last_hidden_state = torch.stack(hidden_states)
+        return last_hidden_state.view(-1, 1, OUTPUT_EMB_SIZE)
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def device(self):
+        return self.model.device
+
+
+class WhisperAudioModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = "openai/whisper-small",
+        num_projector_layers: int = 2,
+        num_tokens_output: int = 10,
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = WhisperAudioModule(model_name_or_path=self.model_name_or_path)
+        self.num_projector_layers = num_projector_layers
+        self.num_tokens_output = num_tokens_output
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        return build_mlp_vector_projector(
+            input_hidden_size=OUTPUT_EMB_SIZE,
+            lm_hidden_size=lm_hidden_size,
+            num_layers=self.num_projector_layers,
+            num_tokens=self.num_tokens_output,
+        )
+
+    @property
+    def name(self) -> str:
+        return "audio_whisper"
+
+    @property
+    def token(self) -> str:
+        return "<speech>"
+
+    @property
+    def data_key(self) -> str:
+        return "speech_audios"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "WhisperAudioModality":
+        self.module.to(dtype=dtype, device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[torch.Tensor]]:
+        row_values = []
+        for row in rows:
+            audios = []
+            for audio_dict in row[self.data_key]:
+                audio_dict = load_audio(
+                    audio_dict,
+                    target_sampling_rate=self.module.feature_extractor.sampling_rate,
+                )
+                audio_processed = self.module.feature_extractor(
+                    audio_dict["array"],
+                    return_tensors="pt",
+                    sampling_rate=audio_dict["sampling_rate"],
+                ).input_features
+                audios.append(audio_processed)
+            row_values.append(torch.stack(audios) if len(audios) > 0 else None)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        audio_features = []
+        for audio_batch in encoded_values:
+            audio_features.append(self.module.forward(audio_batch))
+        return audio_features

src/sonicverse/multi_token/modalities/base_modality.py

@@ -0,0 +1,48 @@
+from typing import Dict, List, Optional, Any
+from abc import ABC, abstractmethod
+from functools import cached_property
+
+import torch.nn as nn
+import torch
+
+
+class Modality(ABC):
+    @abstractmethod
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        pass
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def token(self) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def data_key(self) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def token_width(self) -> int:
+        pass
+
+    @cached_property
+    def token_idx(self) -> int:
+        hash_ = sum(ord(c) ** i for i, c in enumerate(self.token))
+        return -abs(hash_ % 10_000)
+
+    @abstractmethod
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[Any]]:
+        pass
+
+    @abstractmethod
+    def forward(self, encoded_values: List[Any]) -> List[torch.Tensor]:
+        pass
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "Modality":
+        return self

src/sonicverse/multi_token/modalities/bu__init__.py

@@ -0,0 +1,31 @@
+from multi_token.model_utils import MultiTaskType
+from multi_token.modalities.vision_clip import (
+    CLIPVisionModality,
+    OUTPUT_LAYER as CLIP_POOL_LAYER,
+)
+from multi_token.modalities.imagebind import ImageBindModality
+from multi_token.modalities.document_gte import DocumentGTEModality
+from multi_token.modalities.audio_whisper import WhisperAudioModality
+from multi_token.modalities.audio_clap import CLAPAudioModality
+from multi_token.modalities.video_xclip import XCLIPVideoModality
+from multi_token.modalities.audio_descript import DescriptAudioModality
+from multi_token.modalities.audio_mert import MERTAudioModality
+
+MODALITY_BUILDERS = {
+    "vision_clip": lambda: [CLIPVisionModality()],
+    "vision_clip_pool": lambda: [
+        CLIPVisionModality(feature_layer=CLIP_POOL_LAYER, num_tokens_output=10)
+    ],
+    "audio_whisper": lambda: [
+        WhisperAudioModality(
+            num_tokens_output=10, model_name_or_path="openai/whisper-small"
+        )
+    ],
+    "audio_mert": lambda use_multi_task=MultiTaskType.NO_MULTI_TASK, tasks_config=None :[MERTAudioModality(use_multi_task=use_multi_task, tasks_config=tasks_config, num_tokens_output=60, hidden_dim=32, num_conv_layers = 3, num_mlp_layers = 2)],
+    "audio_clap": lambda use_multi_task=MultiTaskType.NO_MULTI_TASK, tasks_config=None :[CLAPAudioModality(use_multi_task=use_multi_task, tasks_config=tasks_config, num_tokens_output=20)],
+    "audio_descript": lambda: [DescriptAudioModality(num_projector_conv_layers=1, num_projector_mlp_layers=1, num_tokens_output=5, codebooks=12)],
+    "video_xclip": lambda: [XCLIPVideoModality(num_tokens_output=10)],
+    "imagebind": lambda: [ImageBindModality()],
+    "document_gte": lambda: [DocumentGTEModality()],
+    "document_gte_x16": lambda: [DocumentGTEModality(num_tokens_output=32)],
+}

src/sonicverse/multi_token/modalities/document_gte.py

@@ -0,0 +1,144 @@
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+import os
+from functools import cache
+from transformers import AutoTokenizer, AutoModel
+
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import build_mlp_vector_projector
+
+GTE_EMBEDDING_SIZE = 1024
+GTE_CONTEXT_WINDOW = 512
+GTE_DEFAULT_MODEL = "thenlper/gte-large"
+DOCUMENT_GTE_FORCE_CPU = "DOCUMENT_GTE_FORCE_CPU"
+
+
+def average_pool(
+    last_hidden_states: torch.Tensor, attention_mask: torch.Tensor
+) -> torch.Tensor:
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+
+@cache
+def _get_tokenizer(model_name_or_path: str = GTE_DEFAULT_MODEL):
+    return AutoTokenizer.from_pretrained(model_name_or_path)
+
+
+def split_text_into_documents(text: str) -> List[str]:
+    from nltk.tokenize import sent_tokenize
+
+    tokenizer = _get_tokenizer(GTE_DEFAULT_MODEL)
+
+    sentences = sent_tokenize(text)
+    documents = [[]]
+
+    for sentence in sentences:
+        sentence_tokens = tokenizer.encode(sentence, add_special_tokens=False)
+        if len(documents[-1]) + len(sentence_tokens) > GTE_CONTEXT_WINDOW:
+            documents.append([])
+        documents[-1].extend(sentence_tokens)
+
+    return [tokenizer.decode(doc) for doc in documents]
+
+
+class DocumentGTEModule(nn.Module):
+    def __init__(self, model_name_or_path: str):
+        super().__init__()
+        self.feature_layer = -2
+        self.model_name_or_path = model_name_or_path
+
+        self.model = AutoModel.from_pretrained("thenlper/gte-large")
+        self.model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, batch_dict) -> torch.Tensor:
+        outputs = self.model(**batch_dict)
+        embeddings = average_pool(
+            outputs.last_hidden_state, batch_dict["attention_mask"]
+        )
+        return embeddings
+
+    @property
+    def embedding_size(self):
+        return GTE_EMBEDDING_SIZE
+
+
+class DocumentGTEModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = GTE_DEFAULT_MODEL,
+        num_projector_layers: int = 2,
+        num_tokens_output: int = 4,
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = DocumentGTEModule(model_name_or_path=self.model_name_or_path)
+        self.tokenizer = _get_tokenizer(model_name_or_path)
+        self.num_projector_layers = num_projector_layers
+        self.num_tokens_output = num_tokens_output
+        self.dtype = torch.float32
+        self.device = "cpu"
+        self.document_gte_device = "cpu"
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        return build_mlp_vector_projector(
+            input_hidden_size=self.module.embedding_size,
+            lm_hidden_size=lm_hidden_size,
+            num_layers=self.num_projector_layers,
+            num_tokens=self.num_tokens_output,
+        )
+
+    @property
+    def name(self) -> str:
+        return "document_gte"
+
+    @property
+    def token(self) -> str:
+        return "<document>"
+
+    @property
+    def data_key(self) -> str:
+        return "documents"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "DocumentGTEModality":
+        self.dtype = dtype
+        self.device = device
+        if DOCUMENT_GTE_FORCE_CPU not in os.environ:
+            # running out of VRAM on 24GB GPU
+            self.document_gte_device = device
+        self.module.to(device=self.document_gte_device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Dict]:
+        row_values = []
+        for row in rows:
+            documents = []
+            for doc in row[self.data_key]:
+                documents.append(doc)
+            documents_tokenized = self.tokenizer(
+                documents,
+                max_length=GTE_CONTEXT_WINDOW,
+                padding=True,
+                truncation=True,
+                return_tensors="pt",
+            )
+            row_values.append(documents_tokenized)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[Dict]) -> List[torch.Tensor]:
+        outputs = []
+        for val in encoded_values:
+            outputs.append(
+                self.module.forward(val.to(device=self.document_gte_device))
+                .to(device=self.device, dtype=self.dtype)
+                .view(-1, 1, self.module.embedding_size)
+            )
+        # batch_size x num_items x 1 x embedding_size
+        return outputs

src/sonicverse/multi_token/modalities/imagebind.py

@@ -0,0 +1,153 @@
+from typing import Dict, List
+import os
+
+import torch
+import torch.nn as nn
+
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import build_mlp_vector_projector
+from multi_token.data_tools import with_local_files
+
+IMAGE_BIND_FORCE_CPU = "IMAGE_BIND_FORCE_CPU"
+IMAGE_BIND_EMBEDDING_SIZE = 1024
+
+
+class ImageBindModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        from imagebind.models import imagebind_model
+        from imagebind import data
+
+        data.BPE_PATH = os.path.join(
+            os.path.dirname(data.__file__), "..", "bpe", "bpe_simple_vocab_16e6.txt.gz"
+        )
+        self.model = imagebind_model.imagebind_huge(pretrained=True)
+        self.model.eval()
+        self.model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, items: Dict) -> torch.Tensor:
+        forward_outs = self.model(items)
+        return forward_outs
+
+    @property
+    def embedding_size(self):
+        return IMAGE_BIND_EMBEDDING_SIZE
+
+
+class ImageBindModality(Modality):
+    def __init__(
+        self,
+        num_projector_layers: int = 2,
+        num_tokens: int = 4,
+        preprocess_device: str = "cpu",
+    ):
+        self.module = ImageBindModule()
+        self.dtype = torch.float32
+        self.device = "cpu"  # used for outputs
+        self.imagebind_device = "cpu"  # used for imagebind model itself
+        self.preprocess_device = preprocess_device  # used for preprocessing
+        self.num_projector_layers = num_projector_layers
+        self.num_tokens = num_tokens
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        return build_mlp_vector_projector(
+            self.module.embedding_size,
+            lm_hidden_size,
+            num_layers=self.num_projector_layers,
+            num_tokens=self.num_tokens,
+        )
+
+    @property
+    def name(self) -> str:
+        return "imagebind"
+
+    @property
+    def token(self) -> str:
+        return "<imagebind>"
+
+    @property
+    def data_key(self) -> str:
+        return "imagebinds"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "ImageBindModality":
+        # we ignore dtype and sometimes device as well
+        self.device = device
+        self.dtype = dtype
+        if IMAGE_BIND_FORCE_CPU not in os.environ:
+            # running out of VRAM on 24GB GPU
+            self.module.to(device=device)
+            self.imagebind_device = device
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[List[Dict]]:
+        from imagebind.models.imagebind_model import ModalityType
+        from imagebind import data
+
+        row_values = []
+        for row in rows:
+            items = []
+            with with_local_files(row[self.data_key]) as item_paths:
+                for item_path in item_paths:
+                    ib_modality = filename_to_imagebind_modality(item_path)
+                    if ib_modality == ModalityType.TEXT:
+                        items.append(
+                            {
+                                ModalityType.TEXT: data.load_and_transform_text(
+                                    [item_path], self.preprocess_device
+                                )
+                            }
+                        )
+                    elif ib_modality == ModalityType.VISION:
+                        items.append(
+                            {
+                                ModalityType.VISION: data.load_and_transform_vision_data(
+                                    [item_path], self.preprocess_device
+                                )
+                            }
+                        )
+                    elif ib_modality == ModalityType.AUDIO:
+                        items.append(
+                            {
+                                ModalityType.AUDIO: data.load_and_transform_audio_data(
+                                    [item_path], self.preprocess_device
+                                )
+                            }
+                        )
+                    else:
+                        raise ValueError(f"Unknown modality type: {ib_modality}")
+            row_values.append(items)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[List[Dict]]) -> List[torch.Tensor]:
+        item_features = []
+        for item_batch in encoded_values:
+            item_batch_emb = []
+            for item in item_batch:
+                item = {
+                    k: v.to(device=self.imagebind_device, dtype=torch.float32)
+                    for k, v in item.items()
+                }
+                item_batch_emb.extend(list(self.module.forward(item).values()))
+            item_features.append(
+                torch.stack(item_batch_emb).to(device=self.device, dtype=self.dtype)
+            )
+        # batch_size x num_items x 1 x embedding_size
+        return item_features
+
+
+def filename_to_imagebind_modality(fn: str) -> str:
+    from imagebind.models.imagebind_model import ModalityType
+
+    _, ext = os.path.splitext(fn)
+    if ext in {".wav"}:
+        return ModalityType.AUDIO
+    elif ext in {".jpg", ".png", ".jpeg"}:
+        return ModalityType.VISION
+    else:
+        return ModalityType.TEXT

src/sonicverse/multi_token/modalities/multi_task_projector_shared.py

@@ -0,0 +1,321 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+import torch.nn.functional as F
+from typing import Dict
+import numpy as np
+
+class CNN(nn.Module):
+    def __init__(self, input_channels = 25, num_class=15):
+        super(CNN, self).__init__()
+        self.aggregator = nn.Parameter(torch.randn((input_channels, 1,1), dtype=torch.float))
+        self.input_channels = input_channels
+
+        # init bn
+        self.bn_init = nn.BatchNorm2d(1)
+
+        # layer 1
+        self.conv_1 = nn.Conv2d(1, 64, 3, padding=1)
+        self.bn_1 = nn.BatchNorm2d(64)
+        self.mp_1 = nn.MaxPool2d((2, 4))
+
+        # layer 2
+        self.conv_2 = nn.Conv2d(64, 128, 3, padding=1)
+        self.bn_2 = nn.BatchNorm2d(128)
+        self.mp_2 = nn.MaxPool2d((2, 4))
+
+        # layer 3
+        self.conv_3 = nn.Conv2d(128, 128, 3, padding=1)
+        self.bn_3 = nn.BatchNorm2d(128)
+        self.mp_3 = nn.MaxPool2d((2, 4))
+
+        # layer 4
+        self.conv_4 = nn.Conv2d(128, 128, 3, padding=1)
+        self.bn_4 = nn.BatchNorm2d(128)
+        self.mp_4 = nn.MaxPool2d((3, 5))
+
+        # layer 5
+        self.conv_5 = nn.Conv2d(128, 64, 3, padding=1)
+        self.bn_5 = nn.BatchNorm2d(64)
+        self.mp_5 = nn.MaxPool2d((3, 3))
+
+        # classifier
+        self.dense = nn.Linear(640, num_class)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, x):
+        aggregator_weights = F.softmax(self.aggregator)
+        # aggregator_weights = aggregator_weights.view(self.input_channels, 1)
+        # print("0 x shape : ")
+        x = (x * aggregator_weights).sum(dim=0)
+
+        # print("aggregator_output shape ", x.shape)
+
+        x = x.unsqueeze(0).unsqueeze(0)
+
+        # print("1 x shape ", x.shape)
+        # init bn
+        x = self.bn_init(x)
+        # print("2 x shape ", x.shape)
+
+        # layer 1
+        x = self.mp_1(nn.ELU()(self.bn_1(self.conv_1(x))))
+        # print("3 x shape ", x.shape)
+
+        # layer 2
+        x = self.mp_2(nn.ELU()(self.bn_2(self.conv_2(x))))
+        # print("4 x shape ", x.shape)
+
+        # layer 3
+        x = self.mp_3(nn.ELU()(self.bn_3(self.conv_3(x))))
+        # print("5 x shape ", x.shape)
+
+        # layer 4
+        x = self.mp_4(nn.ELU()(self.bn_4(self.conv_4(x))))
+        # print("6 x shape ", x.shape)
+
+        # layer 5
+        x = self.mp_5(nn.ELU()(self.bn_5(self.conv_5(x))))
+        # print("7 x shape ", x.shape)
+
+        # classifier
+        x = x.view(x.size(0), -1)
+        # print("8 x shape ", x.shape)
+        x = self.dropout(x)
+        # print("9 x shape ", x.shape)
+        logit = nn.Sigmoid()(self.dense(x))
+        # print("logit shape ", logit.shape)
+
+        return logit
+
+
+class MLP(nn.Module):
+    def __init__(self, input_channels=25, num_class=15):
+        super(MLP, self).__init__()
+        self.aggregator = nn.Parameter(torch.randn((input_channels, 1,1), dtype=torch.float))
+        self.input_channels = input_channels
+
+        self.hidden_layer_1 = nn.Linear(768, 512)
+        self.output = nn.Linear(512, num_class)
+        self.dropout = nn.Dropout(p=0.2)
+        self.loss = self.get_loss() # can return a dict of losses
+
+    def forward(self, x):
+        """
+        x: (B, L, T, H)
+        T=#chunks, can be 1 or several chunks
+        """
+
+        weights = F.softmax(self.aggregator, dim=1)
+        x = (x * weights).sum(dim=1)
+
+        x = x.mean(-2)
+
+        x = self.hidden_layer_1(x)
+        x = F.relu(x)
+        x = self.dropout(x)
+
+        return self.output(x)
+
+    def get_loss(self):
+        return nn.BCEWithLogitsLoss()
+
+class MLPBackbone(nn.Module):
+    def __init__(self, input_features=768, hidden_dim=512):
+        super(MLPBackbone, self).__init__()
+
+        self.hidden_layer_1 = nn.Linear(input_features, hidden_dim)
+        self.dropout = nn.Dropout(p=0.2)
+        self.loss = self.get_loss() # can return a dict of losses
+
+    def forward(self, x):
+        """
+        x: (B, L, T, H)
+        T=#chunks, can be 1 or several chunks
+        """
+
+        x = self.hidden_layer_1(x)
+        x = F.relu(x)
+        x = self.dropout(x)
+
+        return x
+
+    def get_loss(self):
+        return nn.BCEWithLogitsLoss()
+
+class MLPShared(nn.Module):
+    def __init__(self, input_channels=25, num_class=15):
+        super(MLPShared, self).__init__()
+        self.aggregator = nn.Parameter(torch.randn((input_channels, 1,1), dtype=torch.float))
+        self.input_channels = input_channels
+
+        self.hidden_layer_1 = nn.Linear(512, 256)
+        self.output = nn.Linear(256, num_class)
+        self.dropout = nn.Dropout(p=0.2)
+        self.loss = self.get_loss() # can return a dict of losses
+
+    def forward(self, x):
+        """
+        x: (B, L, T, H)
+        T=#chunks, can be 1 or several chunks
+        """
+
+        weights = F.softmax(self.aggregator, dim=1)
+        x = (x * weights).sum(dim=1)
+
+        x = x.mean(-2)
+
+        x = self.hidden_layer_1(x)
+        x = F.relu(x)
+        x = self.dropout(x)
+
+        return self.output(x)
+
+    def get_loss(self):
+        return nn.BCEWithLogitsLoss()
+
+class MLPAggTaskHead(nn.Module):
+    def __init__(self, input_channels: int, input_size: int, output_size: int, use_aggregator: bool, use_time_average: bool, use_sigmoid: bool, use_transpose: bool, num_layers: int, hidden_dim: int, width: int):
+        super(MLPAggTaskHead, self).__init__()
+        if use_aggregator:
+            self.aggregator = nn.Parameter(torch.randn((input_channels), dtype=torch.float))
+        self.use_aggregator = use_aggregator
+        self.use_time_average = use_time_average
+        self.use_transpose = use_transpose
+        self.use_sigmoid = use_sigmoid
+        self.input_channels = input_channels
+        self.output_size = output_size
+        self.width = width
+
+        if self.width > 1:
+            self.layers = nn.ModuleList()
+            for i in range(self.width):
+                mlp_layers = [nn.GELU()]
+                mlp_layers += self._create_mlp_layers(input_size, output_size, num_layers, hidden_dim)
+                if self.use_sigmoid: mlp_layers += [nn.Sigmoid()]
+                self.layers.append(nn.Sequential(*mlp_layers))
+        else:
+            mlp_layers = [nn.GELU()]
+            mlp_layers += self._create_mlp_layers(input_size, output_size, num_layers, hidden_dim)
+            if self.use_sigmoid: mlp_layers += [nn.Sigmoid()]
+            self.layers = nn.Sequential(*mlp_layers)
+
+    def _create_mlp_layers(self, input_size, output_size, num_layers, hidden_dim):
+        if num_layers >=2:
+            layers = [nn.Linear(input_size, hidden_dim)]
+            layers.append(nn.GELU())
+            if num_layers > 2:
+                for _ in range(1, num_layers - 2):
+                    layers += [
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.GELU()
+                    ]
+            layers.append(nn.Linear(hidden_dim, output_size))
+        else:
+            layers = [nn.Linear(input_size, output_size)]
+        return layers
+
+
+    def forward(self, x):
+        if self.use_transpose:
+            x = x.transpose(1, 0)
+        if self.use_time_average:
+            x = x.mean(-2)
+        if self.use_aggregator:
+            aggregator_weights = F.softmax(self.aggregator)
+            aggregator_weights = aggregator_weights.view(self.input_channels, 1)
+            aggregator_output = (x * aggregator_weights).sum(dim=0)
+            aggregator_output = aggregator_output.unsqueeze(dim=0)
+            # print("Agg output ", aggregator_output.shape)
+        else:
+            aggregator_output = x
+
+        if self.width > 1:
+            if (self.input_channels < 1):
+                return torch.cat([layer(aggregator_output.unsqueeze(dim=0)) for layer in self.layers], dim=-2)
+            else:
+                return torch.cat([layer(aggregator_output.unsqueeze(dim=0)).squeeze(dim=0) for layer in self.layers], dim=-2)
+        else:
+            if (self.input_channels < 1):
+                return self.layers(aggregator_output.unsqueeze(dim=0))
+            else:
+                return self.layers(aggregator_output.unsqueeze(dim=0)).squeeze()
+
+
+class MultiTaskModel(nn.Module):
+    def __init__(self, tasks: Dict):
+        super(MultiTaskModel, self).__init__()
+        self.tasks = tasks
+        for task_name, task_head in self.tasks["task_heads"].items():
+            setattr(self, task_name, MLP(13, task_head["output_size"]))
+            if task_name in self.tasks["task_projectors"].keys():
+                task_projector = tasks["task_projectors"][task_name]
+                setattr(self, task_name + "_projector", MLPAggTaskHead(task_projector["input_channels"], task_projector["input_size"], task_projector["output_size"], task_projector["use_aggregator"], task_projector["use_time_average"], task_projector["use_sigmoid"], task_projector["use_transpose"], task_projector["num_layers"], task_projector["hidden_size"], task_projector["width"]))
+
+    def forward(self, x):
+        task_head_outputs = {}
+        task_projector_outputs = []
+
+        backbone_output = x
+
+        for task_name in self.tasks["task_heads"]:
+            if task_name != "lmm_projector":
+                task_head_outputs[task_name] = getattr(self, task_name)(backbone_output)
+                if task_name in self.tasks["task_projectors"].keys():
+                    task_projector_outputs.append(getattr(self, task_name + "_projector")(task_head_outputs[task_name]))
+            else:
+                task_projector_outputs.append(getattr(self, task_name)(backbone_output))
+
+        if len(task_projector_outputs) > 0:
+            task_projector_outputs_unsqueezed = [task_projector_output.unsqueeze(0) for task_projector_output in task_projector_outputs]
+            task_head_outputs["projectors"] = torch.cat(task_projector_outputs_unsqueezed, dim=-2)
+
+        return task_head_outputs
+
+class MultiTaskSharedModel(nn.Module):
+    def __init__(self, tasks: Dict):
+        super(MultiTaskSharedModel, self).__init__()
+        self.tasks = tasks
+        self.use_backbone = False
+        if "backbone" in self.tasks.keys():
+            self.use_backbone = True
+        if self.use_backbone: self.backbone = MLPBackbone(768, 512)
+        for task_name, task_head in self.tasks["task_heads"].items():
+            if task_name != "lmm_projector":
+                setattr(self, task_name, MLPShared(13, task_head["output_size"]))
+            else:
+                setattr(self, task_name, MLPAggTaskHead(task_head["input_channels"], task_head["input_size"], task_head["output_size"], task_head["use_aggregator"], task_head["use_time_average"], task_head["use_sigmoid"], task_head["use_transpose"], task_head["num_layers"], task_head["hidden_size"], task_head["width"]))
+            if task_name in self.tasks["task_projectors"].keys():
+                task_projector = tasks["task_projectors"][task_name]
+                setattr(self, task_name + "_projector", MLPAggTaskHead(task_projector["input_channels"], task_projector["input_size"], task_projector["output_size"], task_projector["use_aggregator"], task_projector["use_time_average"], task_projector["use_sigmoid"], task_projector["use_transpose"], task_projector["num_layers"], task_projector["hidden_size"], task_projector["width"]))
+
+    def forward(self, x):
+        task_head_outputs = {}
+        task_projector_outputs = []
+
+        if self.use_backbone:
+            backbone_output = self.backbone(x)
+        else:
+            backbone_output = x
+
+        #print("Output shape ", backbone_output.shape)
+        for task_name in self.tasks["task_heads"]:
+            #print("task namee ", task_name)
+            if task_name != "lmm_projector":
+                task_head_outputs[task_name] = getattr(self, task_name)(backbone_output)
+                if task_name in self.tasks["task_projectors"].keys():
+                    task_projector_outputs.append(getattr(self, task_name + "_projector")(task_head_outputs[task_name]))
+            else:
+                llm_input = x
+                if self.tasks["task_heads"][task_name]["use_backbone_output"]:
+                    llm_input = backbone_output
+                task_projector_outputs.append(getattr(self, task_name)(llm_input))
+
+        if len(task_projector_outputs) > 0:
+            task_projector_outputs_unsqueezed = [task_projector_output.unsqueeze(0) for task_projector_output in task_projector_outputs]
+            task_head_outputs["projectors"] = torch.cat(task_projector_outputs_unsqueezed, dim=-2)
+
+        return task_head_outputs
+
+
+

src/sonicverse/multi_token/modalities/projectors.py

@@ -0,0 +1,416 @@
+import torch.nn as nn
+import torch
+from typing import Dict
+import numpy as np
+
+import torch.nn.functional as F
+
+def build_patch_mlp_projector(
+    input_hidden_size: int, lm_hidden_size: int, num_layers: int
+) -> nn.Module:
+    modules = [nn.Linear(input_hidden_size, lm_hidden_size)]
+    for _ in range(1, num_layers):
+        modules.append(nn.GELU())
+        modules.append(nn.Linear(lm_hidden_size, lm_hidden_size))
+    return nn.Sequential(*modules)
+
+
+class _MLPVectorProjector(nn.Module):
+    def __init__(
+        self, input_hidden_size: int, lm_hidden_size: int, num_layers: int, width: int
+    ):
+        super(_MLPVectorProjector, self).__init__()
+        self.mlps = nn.ModuleList()
+        for _ in range(width):
+            mlp = [nn.Linear(input_hidden_size, lm_hidden_size)]
+            for _ in range(1, num_layers):
+                mlp.append(nn.GELU())
+                mlp.append(nn.Linear(lm_hidden_size, lm_hidden_size))
+            self.mlps.append(nn.Sequential(*mlp))
+
+    def forward(self, x):
+        output = torch.cat([mlp(x) for mlp in self.mlps], dim=-2)
+        return output
+
+def build_mlp_vector_projector(
+    input_hidden_size: int, lm_hidden_size: int, num_layers: int, num_tokens: int
+):
+    return _MLPVectorProjector(
+        input_hidden_size, lm_hidden_size, num_layers, num_tokens
+    )
+
+class MLPBackbone(nn.Module):
+    def __init__(self, input_size: int, output_size: int, num_layers: int, hidden_dim: int):
+        super(MLPBackbone, self).__init__()
+        self.output_size = output_size
+        mlp_layers = self._create_mlp_layers(input_size, output_size, num_layers, hidden_dim)
+        self.layers = nn.Sequential(*mlp_layers)       
+
+    def _create_conv_layers(self, input_channels, num_conv_layers, hidden_dim):
+        layers = []
+        for _ in range(num_conv_layers):
+            layers += [
+                nn.Conv1d(input_channels, hidden_dim, kernel_size=3, padding=1),
+                nn.GELU(),
+                nn.MaxPool1d(kernel_size=2, stride=2)
+            ]
+            input_channels = hidden_dim
+        return layers
+
+    def _create_mlp_layers(self, input_size, output_size, num_layers, hidden_dim):
+        if num_layers >=2:
+            layers = [nn.Linear(input_size, hidden_dim)]
+            layers.append(nn.GELU())
+            if num_layers > 2:
+                for _ in range(1, num_layers - 2):
+                    layers += [
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.GELU()
+                    ]
+            layers.append(nn.Linear(hidden_dim, output_size))
+        else:
+            layers = [nn.Linear(input_size, output_size)]
+        return layers
+
+    def forward(self, x):
+        return self.layers(x)
+
+class MLPTaskHead(nn.Module):
+    def __init__(self, backbone: nn.Module, input_size: int, output_size: int, num_layers: int, hidden_dim: int, width: int = 1):
+        super(MLPTaskHead, self).__init__()
+        self.backbone = backbone
+        self.width = width
+        if width > 1:
+            self.layers = nn.ModuleList()
+            for i in range(width):
+                mlp_layers = [nn.GELU()]
+                mlp_layers += self._create_mlp_layers(input_size, output_size, num_layers, hidden_dim)
+                self.layers.append(nn.Sequential(*mlp_layers))
+        else:
+            mlp_layers = [nn.GELU()]
+            mlp_layers += self._create_mlp_layers(input_size, output_size, num_layers, hidden_dim)
+            self.layers = nn.Sequential(*mlp_layers)
+
+    def _create_mlp_layers(self, input_size, output_size, num_layers, hidden_dim):
+        if num_layers >=2:
+            layers = [nn.Linear(input_size, hidden_dim)]
+            layers.append(nn.GELU())
+            if num_layers > 2:
+                for _ in range(1, num_layers - 2):
+                    layers += [
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.GELU()
+                    ]
+            layers.append(nn.Linear(hidden_dim, output_size))
+        else:
+            layers = [nn.Linear(input_size, output_size)]
+        return layers
+
+    def _create_conv_layers(self, input_channels, num_conv_layers, hidden_dim):
+        layers = []
+        for _ in range(num_conv_layers):
+            layers += [
+                nn.Conv2d(in_channels = input_channels, out_channels = hidden_dim, kernel_size=(3,3), stride=1, padding=1),
+                nn.GELU(),
+                nn.MaxPool1d(kernel_size=2, stride=2)
+            ]
+            input_channels = hidden_dim
+        return layers
+
+    def forward(self, x):
+        output = self.backbone.forward(x)
+        if self.width > 1:
+            return torch.cat([layer(output) for layer in self.layers], dim=-2)
+        else:
+            return self.layers(output)
+
+class MLPTaskModule(nn.Module):
+    def __init__(self, input_size: int, output_size: int, num_layers: int, hidden_dim: int, width: int = 1):
+        super(MLPTaskModule, self).__init__()
+        self.width = width
+        if width > 1:
+            self.layers = nn.ModuleList()
+            for i in range(width):
+                mlp_layers = [nn.GELU()]
+                mlp_layers += self._create_mlp_layers(input_size, output_size, num_layers, hidden_dim)
+                self.layers.append(nn.Sequential(*mlp_layers))
+        else:
+            mlp_layers = [nn.GELU()]
+            mlp_layers += self._create_mlp_layers(input_size, output_size, num_layers, hidden_dim)
+            self.layers = nn.Sequential(*mlp_layers)
+
+    def _create_mlp_layers(self, input_size, output_size, num_layers, hidden_dim):
+        if num_layers >=2:
+            layers = [nn.Linear(input_size, hidden_dim)]
+            layers.append(nn.GELU())
+            if num_layers > 2:
+                for _ in range(1, num_layers - 2):
+                    layers += [
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.GELU()
+                    ]
+            layers.append(nn.Linear(hidden_dim, output_size))
+        else:
+            layers = [nn.Linear(input_size, output_size)]
+        return layers
+
+    def _create_conv_layers(self, input_channels, num_conv_layers, hidden_dim):
+        layers = []
+        for _ in range(num_conv_layers):
+            layers += [
+                nn.Conv2d(in_channels = input_channels, out_channels = hidden_dim, kernel_size=(3,3), stride=1, padding=1),
+                nn.GELU(),
+                nn.MaxPool1d(kernel_size=2, stride=2)
+            ]
+            input_channels = hidden_dim
+        return layers
+
+    def forward(self, x):
+        if self.width > 1:
+            return torch.cat([layer(x) for layer in self.layers], dim=-2)
+        else:
+            return self.layers(x)
+
+
+class MultiTaskModel(nn.Module):
+    def __init__(self, input_hidden_size: int, input_channels: int, time_average: bool, time_dimension: int, use_aggregator: bool, tasks: Dict):
+        super(MultiTaskModel, self).__init__()
+        self.tasks = tasks
+        self.time_average = time_average
+        self.time_dimension = time_dimension
+        self.use_aggregator = use_aggregator
+        if self.use_aggregator:
+            if (time_average):
+                self.aggregator = nn.Parameter(torch.randn((input_channels, 1), dtype = torch.float))
+            else:
+                self.aggregator = nn.Parameter(torch.randn((input_channels, 1, 1), dtype = torch.float))
+
+        self.backbone = MLPBackbone(input_hidden_size, self.tasks["backbone"]["output_size"], self.tasks["backbone"]["num_layers"], self.tasks["backbone"]["hidden_size"])
+        for task_name, task_head in self.tasks["task_heads"].items():
+            setattr(self, task_name, MLPTaskModule(self.tasks["backbone"]["output_size"], task_head["output_size"], task_head["num_layers"], task_head["hidden_size"], task_head["width"]))
+            if task_name in self.tasks["task_projectors"].keys():
+                task_projector = tasks["task_projectors"][task_name]
+                setattr(self, task_name + "_projector", MLPTaskModule(task_head["output_size"], task_projector["output_size"], task_projector["num_layers"], task_projector["hidden_size"], task_projector["width"]))
+
+    def forward(self, x):
+        task_head_outputs = {}
+        task_projector_outputs = []
+
+        if self.time_average:
+            x = x.mean(self.time_dimension)
+        if self.use_aggregator:
+            aggregator_weights = F.softmax(self.aggregator, dim=0)
+            aggregator_output = (x * aggregator_weights).sum(dim=0)
+            aggregator_output = aggregator_output.unsqueeze(0)
+        else:
+            aggregator_output = x
+
+        backbone_output = self.backbone(aggregator_output)
+
+        for task_name in self.tasks["task_heads"]:
+            if task_name != "lmm_projector":
+                task_head_output = getattr(self, task_name)(backbone_output)
+                min_val = torch.min(task_head_output)
+                max_val = torch.max(task_head_output)
+
+                normalized_task_head_output = (task_head_output - min_val) / (max_val - min_val)
+                task_head_outputs[task_name] = normalized_task_head_output
+                if task_name in self.tasks["task_projectors"].keys():
+                    task_projector_outputs.append(getattr(self, task_name + "_projector")(task_head_outputs[task_name]))
+            else:
+                task_projector_outputs.append(getattr(self, task_name)(backbone_output))
+
+        task_projector_outputs_unsqueezed = [task_projector_output.unsqueeze(0) for task_projector_output in task_projector_outputs]
+        if len(task_projector_outputs_unsqueezed) > 0:
+            task_head_outputs["projectors"] = torch.cat(task_projector_outputs_unsqueezed, dim=-2)
+
+        return task_head_outputs
+
+
+def build_mt_vector_projector(
+        input_hidden_size: int, lm_hidden_size: int, tasks: Dict
+):
+    projector = nn.ModuleDict()
+    projector["backbone"] = MLPBackbone(input_hidden_size, tasks["backbone"]["output_size"], tasks["backbone"]["num_layers"], tasks["backbone"]["hidden_size"])
+    for task_name, task_head in tasks["task_heads"].items():
+        projector[task_name] = MLPTaskHead(projector["backbone"], task_head["hidden_size"], task_head["output_size"], task_head["num_layers"], task_head["hidden_size"], task_head["width"])
+
+    return projector
+
+class Attention(nn.Module):
+    def __init__(self, input_dim, hidden_dim):
+        super(Attention, self).__init__()
+        self.linear_in = nn.Linear(input_dim, hidden_dim)
+        self.linear_out = nn.Linear(hidden_dim, 1)
+
+    def forward(self, x):
+        # Input shape: (batch_size, seq_len, input_dim)
+        energy = torch.tanh(self.linear_in(x))
+        attention_scores = torch.softmax(self.linear_out(energy), dim=1)
+        context_vector = torch.sum(attention_scores * x, dim=1)
+        return context_vector
+
+class _CNNAttentionTokenizer(nn.Module):
+    def __init__(self, input_channels, output_size, width, hidden_dim, num_conv_layers):
+        super(_CNNAttentionTokenizer, self).__init__()
+        self.width = width
+        self.cnns = nn.ModuleList()
+        self.attentions = nn.ModuleList()
+        for _ in range(width):
+            cnn = self._create_conv_layers(input_channels, num_conv_layers)
+            self.cnns.append(cnn)
+            attention = [Attention(hidden_dim, 125)]
+            linear_input_size = hidden_dim
+            attention.append(nn.Linear(linear_input_size, output_size))
+            self.attentions.append(nn.Sequential(*attention))
+
+
+    def _create_conv_layers(self, input_channels, num_conv_layers):
+        layers = []
+        in_channels = input_channels
+        for _ in range(num_conv_layers):
+            layers += [
+                nn.Conv1d(in_channels, 64, kernel_size=3, padding=1),
+                nn.ReLU(),
+                nn.MaxPool1d(kernel_size=2, stride=2)
+            ]
+            in_channels = 64
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        outputs = []
+        for token in range(self.width):
+            # Input shape: (batch_size, input_channels, sequence_length)
+            token_output = self.cnns[token](x)  # Apply convolutional layers
+            token_output = token_output.permute(0, 2, 1)  # Reshape for attention mechanism (batch_size, sequence_length, input_dim
+            token_output = self.attentions[token](token_output)  # Apply attention mechanism
+            outputs.append(token_output)
+        output = torch.cat(outputs, dim=-2)
+        output = torch.stack([output])
+        return output
+
+def build_attentive_cnn_projector(
+    input_channels: int, lm_hidden_size: int, num_tokens: int, hidden_dim: int, num_layers: int
+    ):
+    return _CNNAttentionTokenizer(input_channels, lm_hidden_size, num_tokens, hidden_dim, num_layers)
+
+class _CNNMLPProjector(nn.Module):
+    def __init__(self, input_channels, input_size, output_size = 4096, width = 5, hidden_dim = 64, num_conv_layers = 1, num_mlp_layers = 2):
+        super(_CNNMLPProjector, self).__init__()
+        self.width = width
+        self.cnnmlps = nn.ModuleList()
+        for _ in range(self.width):
+            cnnmlp = self._create_conv_layers(input_channels, num_conv_layers, hidden_dim)
+            cnnmlp.append(nn.Flatten())
+            cnn_output_size = hidden_dim*((input_size + 2*1 - 3*num_conv_layers) // (2**num_conv_layers) + 1)
+            cnnmlp.append(nn.Linear(cnn_output_size, output_size))
+            cnnmlp.append(nn.GELU())
+            cnnmlp += self._create_mlp_layers(output_size, output_size, num_mlp_layers, output_size)
+            self.cnnmlps.append(nn.Sequential(*cnnmlp))
+
+    def _create_conv_layers(self, input_channels, num_conv_layers, hidden_dim):
+        layers = []
+        for _ in range(num_conv_layers):
+            layers += [
+                nn.Conv1d(input_channels, hidden_dim, kernel_size=3, padding=1),
+                nn.GELU(),
+                nn.MaxPool1d(kernel_size=2, stride=2)
+            ]
+            input_channels = hidden_dim
+        return layers
+
+    def _create_mlp_layers(self, input_size, output_size, num_layers, hidden_dim):
+        if num_layers >=2:
+            layers = [nn.Linear(input_size, hidden_dim)]
+            layers.append(nn.GELU())
+            if num_layers > 2:
+                for _ in range(1, num_layers - 2):
+                    layers += [
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.GELU()
+                    ]
+            layers.append(nn.Linear(hidden_dim, output_size))
+        else:
+            layers = [nn.Linear(input_size, output_size)]
+        return layers
+
+    def forward(self, x):
+        return torch.stack([torch.cat([cnnmlp(x) for cnnmlp in self.cnnmlps], dim=-2)])
+
+def build_cnn_mlp_projector(
+    input_channels: int, input_size: int, lm_hidden_size: int, num_tokens: int, hidden_dim: int, num_conv_layers: int, num_mlp_layers: int
+    ):
+    return _CNNMLPProjector(input_channels, input_size, lm_hidden_size, num_tokens, hidden_dim, num_conv_layers, num_mlp_layers)
+
+class _MultiLayeredCNNMLPProjector(nn.Module):
+    def __init__(self, input_channels, input_size, num_feature_layers, output_size = 4096, width = 5, hidden_dim = 64, num_conv_layers = 1, num_mlp_layers = 2):
+        super(_MultiLayeredCNNMLPProjector, self).__init__()
+        self.width = width
+        self.num_feature_layers = num_feature_layers
+        self.cnnmlps = nn.ModuleList()
+        for _ in range(self.width*self.num_feature_layers):
+            cnnmlp = self._create_conv_layers(input_channels, num_conv_layers, hidden_dim)
+            cnnmlp += [nn.GELU()]
+            cnnmlp += self._create_mlp_layers(input_size, output_size, num_mlp_layers, output_size)
+            self.cnnmlps.append(nn.Sequential(*cnnmlp))
+
+    def _create_conv_layers(self, input_channels, num_conv_layers, hidden_size):
+        layers = []
+
+        if input_channels >= hidden_size:
+            hidden_dim = int(input_channels/2)
+        else:
+            hidden_dim = hidden_size
+
+        layers += [nn.Conv1d(in_channels=input_channels, out_channels=hidden_dim, kernel_size=3, stride=1, padding=1), nn.GELU()]
+        if num_conv_layers > 2:
+            for _ in range(num_conv_layers - 2):
+                if hidden_dim/2 >= hidden_size:
+                    output_dim = int(hidden_dim/2)
+                else:
+                    output_dim = hidden_size
+                layers += [
+                    nn.Conv1d(in_channels=hidden_dim, out_channels=output_dim, kernel_size=3, stride=1, padding=1),
+                    nn.GELU(),
+                ]
+                hidden_dim = output_dim
+        layers += [nn.Conv1d(in_channels=hidden_dim, out_channels=1, kernel_size=3, stride=1, padding=1)]
+        return layers
+
+    def _create_mlp_layers(self, input_size, output_size, num_layers, hidden_dim):
+        if num_layers >=2:
+            layers = [nn.Linear(input_size, hidden_dim)]
+            layers.append(nn.GELU())
+            if num_layers > 2:
+                for _ in range(1, num_layers - 2):
+                    layers += [
+                        nn.Linear(hidden_dim, hidden_dim),
+                        nn.GELU()
+                    ]
+            layers.append(nn.Linear(hidden_dim, output_size))
+        else:
+            layers = [nn.Linear(input_size, output_size)]
+        return layers
+
+    def forward(self, x):
+        print("X SHAPE ", x.shape)
+        inp_feature_layers = []
+        for feature_id in range(self.num_feature_layers):
+            in_feat_layer = x[feature_id].unsqueeze(0).permute(0,2,1)
+            inp_feature_layers.append(in_feat_layer)
+
+        outputs = []
+        for layer_count in range(self.width*self.num_feature_layers):
+            feature_id = int(layer_count/self.width)
+            outputs+=[self.cnnmlps[layer_count](inp_feature_layers[feature_id])]
+        
+        return torch.cat(outputs, dim=-2)
+
+
+def build_multi_layer_cnn_mlp_projector(
+    input_channels: int, input_size: int, num_feature_layers: int, lm_hidden_size: int, num_tokens: int, hidden_dim: int, num_conv_layers: int, num_mlp_layers: int
+    ):
+    assert(num_tokens % num_feature_layers == 0)
+    width = int(num_tokens/num_feature_layers)
+    return _MultiLayeredCNNMLPProjector(input_channels, input_size, num_feature_layers, lm_hidden_size, width, hidden_dim, num_conv_layers, num_mlp_layers)
+

src/sonicverse/multi_token/modalities/video_xclip.py

@@ -0,0 +1,113 @@
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from transformers import AutoProcessor, AutoModel
+
+from multi_token.data_tools import load_video
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_mlp_vector_projector,
+)
+
+
+OUTPUT_EMB_SIZE = 512
+
+
+class XCLIPVideoModule(nn.Module):
+    def __init__(self, model_name_or_path: str):
+        super().__init__()
+        self.model_name_or_path = model_name_or_path
+        self.model = None
+        self.processor = None
+
+        self.load_model()
+
+    def load_model(self):
+        self.model = AutoModel.from_pretrained(self.model_name_or_path)
+        self.processor = AutoProcessor.from_pretrained(self.model_name_or_path)
+        self.model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, video_inputs) -> torch.Tensor:
+        with torch.no_grad():
+            outputs = self.model(**(video_inputs.to(device=self.device)))
+
+        emb = outputs.video_embeds.to(device=self.device, dtype=self.dtype).view(
+            -1, 1, OUTPUT_EMB_SIZE
+        )
+        return emb
+
+    @property
+    def dtype(self):
+        return self.model.dtype
+
+    @property
+    def device(self):
+        return self.model.device
+
+
+class XCLIPVideoModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = "microsoft/xclip-base-patch32",
+        num_projector_layers: int = 2,
+        num_tokens_output: int = 10,
+    ):
+        self.model_name_or_path = model_name_or_path
+        self.module = XCLIPVideoModule(model_name_or_path=self.model_name_or_path)
+        self.num_projector_layers = num_projector_layers
+        self.num_tokens_output = num_tokens_output
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        return build_mlp_vector_projector(
+            input_hidden_size=OUTPUT_EMB_SIZE,
+            lm_hidden_size=lm_hidden_size,
+            num_layers=self.num_projector_layers,
+            num_tokens=self.num_tokens_output,
+        )
+
+    @property
+    def name(self) -> str:
+        return "video_xclip"
+
+    @property
+    def token(self) -> str:
+        return "<video>"
+
+    @property
+    def data_key(self) -> str:
+        return "videos"
+
+    @property
+    def token_width(self) -> int:
+        return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "XCLIPVideoModality":
+        self.module.to(dtype=dtype, device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[Dict]]:
+        row_values = []
+        for row in rows:
+            video_arrays = [
+                load_video(
+                    video_info,
+                )
+                for video_info in row[self.data_key]
+            ]
+            videos_enc = self.module.processor(
+                videos=[list(video) for video in video_arrays],
+                text=["IGNORE"],
+                return_tensors="pt",
+                padding=True,
+            )
+            row_values.append(videos_enc)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        video_features = []
+        for video_batch in encoded_values:
+            video_features.append(self.module.forward(video_batch))
+        return video_features

src/sonicverse/multi_token/modalities/vision_clip.py

@@ -0,0 +1,178 @@
+from typing import Dict, List, Tuple, Optional
+
+import torch
+import torch.nn as nn
+from transformers import CLIPVisionModel, CLIPImageProcessor
+from PIL import Image
+
+from multi_token.modalities.base_modality import Modality
+from multi_token.modalities.projectors import (
+    build_patch_mlp_projector,
+    build_mlp_vector_projector,
+)
+from multi_token.data_tools import load_image
+
+PATCH_LAYER = -2
+OUTPUT_LAYER = -1
+OUTPUT_EMB_SIZE = 1024
+
+
+class CLIPVisionModule(nn.Module):
+    def __init__(self, model_name_or_path: str, feature_layer: int = PATCH_LAYER):
+        super().__init__()
+        self.feature_layer = feature_layer
+        self.model_name_or_path = model_name_or_path
+        self.image_processor = None
+        self.image_model = None
+
+        self.load_model()
+
+    def load_model(self):
+        self.image_processor = CLIPImageProcessor.from_pretrained(
+            self.model_name_or_path
+        )
+        self.image_model = CLIPVisionModel.from_pretrained(self.model_name_or_path)
+        self.image_model.requires_grad_(False)
+
+    @torch.no_grad()
+    def forward(self, images) -> torch.Tensor:
+        if self.feature_layer == PATCH_LAYER:
+            image_forward_outs = self.image_model(
+                images.to(device=self.device, dtype=self.dtype),
+                output_hidden_states=True,
+            )
+            image_features = image_forward_outs.hidden_states[self.feature_layer]
+            image_features = image_features[:, 1:].to(images.dtype)
+        else:
+            image_forward_outs = self.image_model(
+                images.to(device=self.device, dtype=self.dtype),
+            )
+            image_features = image_forward_outs.pooler_output.to(images.dtype).view(
+                -1, 1, OUTPUT_EMB_SIZE
+            )
+        return image_features
+
+    @property
+    def dtype(self):
+        return self.image_model.dtype
+
+    @property
+    def device(self):
+        return self.image_model.device
+
+    @property
+    def config(self):
+        return self.image_model.config
+
+    @property
+    def hidden_size(self):
+        return self.config.hidden_size
+
+    @property
+    def num_patches(self):
+        return (self.config.image_size // self.config.patch_size) ** 2
+
+
+def _expand2square(pil_img: Image, background_color: Tuple) -> Image:
+    width, height = pil_img.size
+    if width == height:
+        return pil_img
+    elif width > height:
+        result = Image.new(pil_img.mode, (width, width), background_color)
+        result.paste(pil_img, (0, (width - height) // 2))
+        return result
+    else:
+        result = Image.new(pil_img.mode, (height, height), background_color)
+        result.paste(pil_img, ((height - width) // 2, 0))
+        return result
+
+
+class CLIPVisionModality(Modality):
+    def __init__(
+        self,
+        model_name_or_path: str = "openai/clip-vit-large-patch14-336",
+        pad_non_square_images: bool = False,
+        num_projector_layers: int = 2,
+        feature_layer: int = PATCH_LAYER,
+        num_tokens_output: Optional[int] = None,
+    ):
+        if feature_layer not in [PATCH_LAYER, OUTPUT_LAYER]:
+            raise ValueError(
+                f"feature_layer must be one of {PATCH_LAYER} or {OUTPUT_LAYER}"
+            )
+        if (feature_layer == PATCH_LAYER) != (num_tokens_output is None):
+            raise ValueError(
+                "num_tokens_output must be None if feature_layer is PATCH_LAYER"
+            )
+        self.model_name_or_path = model_name_or_path
+        self.module = CLIPVisionModule(
+            model_name_or_path=self.model_name_or_path, feature_layer=feature_layer
+        )
+        self.pad_non_square_images = pad_non_square_images
+        self.num_projector_layers = num_projector_layers
+        self.num_tokens_output = num_tokens_output
+
+    def build_projector(self, lm_hidden_size: int) -> nn.Module:
+        if self.module.feature_layer == PATCH_LAYER:
+            return build_patch_mlp_projector(
+                self.module.hidden_size,
+                lm_hidden_size,
+                num_layers=self.num_projector_layers,
+            )
+        else:
+            return build_mlp_vector_projector(
+                input_hidden_size=OUTPUT_EMB_SIZE,
+                lm_hidden_size=lm_hidden_size,
+                num_layers=self.num_projector_layers,
+                num_tokens=self.num_tokens_output,
+            )
+
+    @property
+    def name(self) -> str:
+        return "vision_clip"
+
+    @property
+    def token(self) -> str:
+        return "<image>"
+
+    @property
+    def data_key(self) -> str:
+        return "images"
+
+    @property
+    def token_width(self) -> int:
+        if self.module.feature_layer == PATCH_LAYER:
+            return self.module.num_patches
+        else:
+            return self.num_tokens_output
+
+    def to(self, dtype: torch.dtype, device: torch.device) -> "CLIPVisionModality":
+        self.module.to(dtype=dtype, device=device)
+        return self
+
+    def preprocess_rows(self, rows: List[Dict]) -> List[Optional[torch.Tensor]]:
+        row_values = []
+        for row in rows:
+            images = []
+            for image_fn in row[self.data_key]:
+                image_obj = load_image(image_fn)
+                if self.pad_non_square_images:
+                    image_obj = _expand2square(
+                        image_obj,
+                        tuple(
+                            int(x * 255) for x in self.module.image_processor.image_mean
+                        ),
+                    )
+                image = self.module.image_processor.preprocess(
+                    image_obj, return_tensors="pt"
+                )["pixel_values"][0]
+                images.append(image)
+            row_values.append(torch.stack(images) if len(images) > 0 else None)
+        return row_values
+
+    @torch.no_grad()
+    def forward(self, encoded_values: List[torch.Tensor]) -> List[torch.Tensor]:
+        image_features = []
+        for image_batch in encoded_values:
+            image_features.append(self.module.forward(image_batch))
+        return image_features

src/sonicverse/multi_token/model_utils.py

@@ -0,0 +1,112 @@
+from typing import List, Dict
+import logging
+import torch
+
+from enum import Enum
+
+class MultiTaskType(Enum):
+    NO_MULTI_TASK = 0
+    SIMPLE_MULTI_TASK = 1
+    PROJECTED_MULTI_TASK = 2
+
+def _find_all_linear_names(model) -> List[str]:
+    cls = torch.nn.Linear
+    lora_module_names = set()
+    for name, module in model.named_modules():
+        if isinstance(module, cls):
+            names = name.split(".")
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+
+    if "lm_head" in lora_module_names:
+        lora_module_names.remove("lm_head")
+    return list(lora_module_names)
+
+
+def maybe_zero_3(param, ignore_status=False, name=None):
+    from deepspeed import zero
+    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
+
+    if hasattr(param, "ds_id"):
+        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE:
+            if not ignore_status:
+                logging.warning(
+                    f"{name}: param.ds_status != ZeroParamStatus.NOT_AVAILABLE: {param.ds_status}"
+                )
+        with zero.GatheredParameters([param]):
+            param = param.data.detach().cpu().clone()
+    else:
+        param = param.detach().cpu().clone()
+    return param
+
+
+def get_peft_state(named_params, bias) -> Dict:
+    if bias == "none":
+        to_return = {k: t for k, t in named_params if "lora_" in k}
+    elif bias == "all":
+        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
+    elif bias == "lora_only":
+        to_return = {}
+        maybe_lora_bias = {}
+        lora_bias_names = set()
+        for k, t in named_params:
+            if "lora_" in k:
+                to_return[k] = t
+                bias_name = k.split("lora_")[0] + "bias"
+                lora_bias_names.add(bias_name)
+            elif "bias" in k:
+                maybe_lora_bias[k] = t
+        for k, t in maybe_lora_bias:
+            if bias_name in lora_bias_names:
+                to_return[bias_name] = t
+    else:
+        raise NotImplementedError()
+    to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()}
+    return to_return
+
+
+def get_peft_state_non_lora(named_params, task_names) -> Dict:
+    to_return = {}
+    for k, t in named_params:
+        if "lora_" not in k:
+            task_name_in_k = False
+            for task_name in task_names:
+                if task_name in k:
+                    task_name_in_k = True
+            if t.requires_grad or task_name_in_k:
+                to_return[k] = t
+    to_return = {
+        k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()
+    }
+    return to_return
+
+
+def make_model_lora(model, training_args: "TrainingArguments"):
+    from peft import LoraConfig, get_peft_model
+
+    lora_config = LoraConfig(
+        r=training_args.lora_r,
+        lora_alpha=training_args.lora_alpha,
+        target_modules=_find_all_linear_names(model),
+        lora_dropout=training_args.lora_dropout,
+        bias=training_args.lora_bias,
+        task_type="CAUSAL_LM",
+    )
+    if training_args.bits == 16:
+        if training_args.bf16:
+            model.to(torch.bfloat16)
+        if training_args.fp16:
+            model.to(torch.float16)
+
+    model = get_peft_model(model, lora_config)
+    return model
+
+
+def fix_tokenizer(tokenizer):
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.unk_token
+    if tokenizer.mask_token is None:
+        tokenizer.mask_token = tokenizer.unk_token
+    if tokenizer.cls_token is None:
+        tokenizer.cls_token = tokenizer.unk_token
+    if tokenizer.sep_token is None:
+        tokenizer.sep_token = tokenizer.unk_token

src/sonicverse/multi_token/training.py

@@ -0,0 +1,344 @@
+from typing import Optional, List
+from dataclasses import field, dataclass
+import logging
+import subprocess
+import pathlib
+import torch
+import shutil
+import glob
+import os
+import json
+
+import transformers
+from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
+from transformers import Trainer
+
+from multi_token.training_data import (
+    DataArguments,
+    LMMDataset,
+    DataCollatorForSupervisedLMMDataset,
+)
+from multi_token.model_utils import (
+    make_model_lora,
+    get_peft_state,
+    get_peft_state_non_lora,
+    fix_tokenizer,
+    MultiTaskType
+)
+from multi_token.modalities.base_modality import Modality
+
+
+README_TEMPLATE = """
+---
+license: apache-2.0
+base_model: {base_model}
+dataset: {dataset}
+tags:
+  - finetuned
+  - multimodal
+inference: false
+---
+
+These are weights for a version of `{base_model}` finetuned for multimodal applications. 
+
+### Modalities
+
+{modalities}
+
+### Usage
+
+GitHub: https://github.com/sshh12/multi_token (includes training scripts and basic inference server)
+
+### Dataset
+
+{dataset} ({num_examples} examples)
+
+```
+{dataset_example}
+```
+
+### Training Device(s)
+
+```
+{training_devices_dump}
+```
+
+
+### Model
+
+```
+{repr_model}
+```
+
+"""
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    remove_unused_columns: bool = field(default=False)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=512,
+        metadata={
+            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
+        },
+    )
+    double_quant: bool = field(
+        default=True,
+        metadata={
+            "help": "Compress the quantization statistics through double quantization."
+        },
+    )
+    quant_type: str = field(
+        default="nf4",
+        metadata={
+            "help": "Quantization data type to use. Should be one of `fp4` or `nf4`."
+        },
+    )
+    pretrain_projectors: bool = field(default=False)
+    pretrained_projectors_path: Optional[str] = field(default=None)
+    pretrained_projectors_config: Optional[str] = field(default=None)
+    bits: int = field(default=16, metadata={"help": "How many bits to use."})
+    lora_enable: bool = False
+    lora_r: int = 64
+    lora_alpha: int = 16
+    lora_dropout: float = 0.05
+    lora_weight_path: str = ""
+    lora_bias: str = "none"
+
+
+@dataclass
+class ModelArguments:
+    model_name_or_path: str = field(default="mistralai/Mistral-7B-Instruct-v0.1")
+    model_cls: str = field(default="MistralLMMForCausalLM")
+    modality_builder: str = field(default="vision_clip")
+    use_multi_task: int = field(default=MultiTaskType.PROJECTED_MULTI_TASK)
+    tasks_config: str = field(default="src/sonicverse/configs/tasks.json")
+    model_lora_path: Optional[str] = field(default="annabeth97c/sonicverse")
+
+
+class LMMTrainer(Trainer):
+    def _save_checkpoint(self, model, trial, metrics=None):
+        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
+
+        run_dir = self._get_output_dir(trial=trial)
+        output_dir = os.path.join(run_dir, checkpoint_folder)
+        self._save_extras(output_dir)
+
+        super(LMMTrainer, self)._save_checkpoint(model, trial, metrics)
+
+    def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        self._save_extras(output_dir)
+        super(LMMTrainer, self)._save(output_dir, state_dict)
+        for unused_dir in glob.iglob(os.path.join(output_dir, "global_step*")):
+            shutil.rmtree(unused_dir)
+
+    def _save_extras(self, output_dir: Optional[str] = None):
+        self.model.config.save_pretrained(output_dir)
+
+        task_names = []
+        for m in self.model.modalities:
+            task_names += m.tasks["task_heads"].keys()
+
+        non_lora_state_dict = get_peft_state_non_lora(self.model.named_parameters(), task_names)
+        torch.save(
+            non_lora_state_dict,
+            os.path.join(output_dir, "non_lora_trainables.bin"),
+        )
+
+
+def _get_training_devices_dump() -> str:
+    out = subprocess.check_output(
+        ["nvidia-smi", "--query-gpu=gpu_name,gpu_bus_id,vbios_version", "--format=csv"]
+    )
+    return out.decode("utf-8").strip()
+
+
+def train_for_modalities(
+    model_cls,
+    training_args: TrainingArguments,
+    model_args: ModelArguments,
+    train_data_args: DataArguments,
+    evaluation_data_args: DataArguments,
+    modalities: List[Modality],
+):
+    for m in modalities:
+        m.to(
+            dtype=torch.bfloat16 if training_args.bf16 else torch.float16,
+            device=training_args.device,
+        )
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        model_max_length=training_args.model_max_length,
+        padding_side="right",
+        use_fast=False,
+    )
+    fix_tokenizer(tokenizer)
+
+    train_dataset = LMMDataset(train_data_args, tokenizer, modalities)
+    evaluation_dataset = LMMDataset(evaluation_data_args, tokenizer, modalities)
+    collator = DataCollatorForSupervisedLMMDataset(tokenizer, modalities)
+
+    model = model_cls.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+    )
+    model.to(
+            dtype=torch.bfloat16 if training_args.bf16 else torch.float16,
+            device=training_args.device,
+        )
+    model.modalities = modalities
+    model.config.use_cache = False
+    model.config.model_cls = model_cls.__name__
+    model.config.modality_builder = model_args.modality_builder
+
+    if training_args.gradient_checkpointing:
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
+
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+
+            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+    if model_args.model_lora_path:
+        raise ValueError(
+            "LoRA path not supported for training -- set the output path to an existing model to resume training"
+        )
+
+    if training_args.lora_enable:
+        logging.info("Adding LoRA adapters...")
+        model = make_model_lora(model, training_args)
+
+    if training_args.pretrained_projectors_path:
+        projector_weights_og = torch.load(
+            training_args.pretrained_projectors_path, map_location="cpu"
+        )
+        if model_args.use_multi_task==MultiTaskType.SIMPLE_MULTI_TASK:
+            projector_weights = {}
+            for k, v in projector_weights_og.items():
+                for m in modalities:
+                    for task_name in m.tasks["task_heads"].keys():
+                        if task_name in k:
+                            projector_weights[k] = v
+        else:
+            projector_weights = {
+                k: v for k, v in projector_weights_og.items() if "_lmm_projector" in k
+            }
+
+    elif training_args.pretrained_projectors_config:
+        with open(training_args.pretrained_projectors_config, "r") as f:
+            pretrained_weights_config = json.load(f)
+
+        projector_weights = {}
+
+        for pretrained_path_info in pretrained_weights_config["pretrained_paths"]:
+            pretrained_path = pretrained_path_info["path"]
+            components = pretrained_path_info["components"]
+            use_prefix = pretrained_path_info["use_prefix"]
+            prefix = pretrained_path_info["prefix"]
+            
+            pretrained_weights = torch.load(pretrained_path, map_location="cpu")
+            
+            for k, v in pretrained_weights.items():
+                if any(component in k for component in components):
+                    weight_key = k
+                    if use_prefix:
+                        weight_key = prefix + "." + k
+                    projector_weights[weight_key] = v
+
+    else:
+        projector_weights = {}
+
+    model.get_model().initialize_modules(modalities, projector_weights)
+
+    task_names = []
+    tasks = {}
+    for m in model.modalities:
+        if m.use_multi_task != MultiTaskType.NO_MULTI_TASK:
+            tasks = m.tasks
+            task_names += m.tasks["task_heads"].keys()
+
+    if training_args.pretrain_projectors:
+        model.requires_grad_(False)
+        for m in modalities:
+            if m.use_multi_task == MultiTaskType.SIMPLE_MULTI_TASK:
+                for task_name in m.tasks["task_heads"].keys():
+                    task_model = getattr(model.get_model(), m.name + "_" + task_name)
+                    for p in task_model.parameters():
+                        p.requires_grad = True
+            elif m.use_multi_task == MultiTaskType.PROJECTED_MULTI_TASK:
+                proj = getattr(model.get_model(), m.name + "_lmm_projector")
+                
+                if "backbone" in m.tasks.keys():
+                    backbone = getattr(proj,  "backbone")
+                    for backbone_param in backbone.parameters():
+                        backbone_param.requires_grad = tasks["backbone"]["requires_grad"]
+
+                for task in task_names:
+                    task_head = getattr(proj, task)
+                    for task_head_param in task_head.parameters():
+                        task_head_param.requires_grad = tasks["task_heads"][task]["requires_grad"]
+                    if task in tasks["task_projectors"]:
+                        task_projector = getattr(proj, task + "_projector")
+                        for task_projector_param in task_projector.parameters():
+                            task_projector_param.requires_grad = tasks["task_projectors"][task]["requires_grad"]
+
+            else:
+                proj = getattr(model.get_model(), m.name + "_lmm_projector")
+                for p in proj.parameters():
+                    p.requires_grad = True
+
+    os.makedirs(training_args.output_dir, exist_ok=True)
+    with open(
+        os.path.join(training_args.output_dir, "model_named_parameters.txt"), "w"
+    ) as f:
+        for name, param in model.named_parameters():
+            f.write(f"{name} {param.shape} {param.requires_grad}\n")
+
+    with open(os.path.join(training_args.output_dir, "README.md"), "w") as f:
+        modalities_text = [
+            f"* {m.__class__.__name__} (use `{m.token}` in text and provide `{m.data_key}`, encoded as {m.token_width} tokens)"
+            for m in modalities
+        ]
+        readme_text = README_TEMPLATE.format(
+            base_model=model_args.model_name_or_path,
+            dataset=train_data_args.dataset_path,
+            dataset_example=repr(train_dataset.get_example()),
+            num_examples=len(train_dataset),
+            modalities="\n".join(modalities_text),
+            training_devices_dump=_get_training_devices_dump(),
+            repr_model=f"{model_cls.__name__}.model =\n\n{repr(model)}",
+        )
+        f.write(readme_text)
+
+    trainer = LMMTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        args=training_args,
+        data_collator=collator,
+        train_dataset=train_dataset,
+        eval_dataset=evaluation_dataset,
+    )
+
+    if list(pathlib.Path(training_args.output_dir).glob(f"{PREFIX_CHECKPOINT_DIR}-*")):
+        trainer.train(resume_from_checkpoint=True)
+    else:
+        trainer.train()
+
+    trainer.save_state()
+
+    model.config.use_cache = True
+    model.config.save_pretrained(training_args.output_dir)
+    state_dict = get_peft_state(model.named_parameters(), training_args.lora_bias)
+    model.save_pretrained(training_args.output_dir, state_dict=state_dict)
+
+    non_lora_state_dict = get_peft_state_non_lora(model.named_parameters(), task_names)
+    torch.save(
+        non_lora_state_dict,
+        os.path.join(training_args.output_dir, "non_lora_trainables.bin"),
+    )

src/sonicverse/multi_token/training_data.py

@@ -0,0 +1,133 @@
+from typing import List, Dict, Sequence
+from dataclasses import dataclass, field
+import logging
+import os
+
+from torch.utils.data import Dataset
+from datasets import load_from_disk, load_dataset, Dataset as HFDataset
+import transformers
+import torch
+
+from multi_token.modalities.base_modality import Modality
+from multi_token.constants import IGNORE_INDEX
+from multi_token.data_tools import encode_chat, encode_chat_multitask
+from multi_token.model_utils import MultiTaskType
+
+
+@dataclass
+class DataArguments:
+    dataset_path: str = field(
+        default=None, metadata={"help": "Path to the training data."}
+    )
+
+@dataclass
+class TrainDataArguments:
+    train_dataset_path: str = field(
+        default=None, metadata={"help": "Path to the training data."}
+    )
+
+@dataclass
+class EvaluationDataArguments:
+    evaluation_dataset_path: str = field(
+        default=None, metadata={"help": "Path to the evaluation data."}
+    )
+
+
+def _resolve_dataset(path: str) -> HFDataset:
+    if os.path.exists(path):
+        return load_from_disk(path)
+    else:
+        return load_dataset(path, split="train", data_files="*.arrow")
+
+
+class LMMDataset(Dataset):
+    def __init__(
+        self,
+        data_args: DataArguments,
+        tokenizer: transformers.PreTrainedTokenizer,
+        modalities: List[Modality],
+    ):
+        super(LMMDataset, self).__init__()
+        self.dataset = _resolve_dataset(data_args.dataset_path)
+        self.tokenizer = tokenizer
+        self.modalities = modalities
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def get_example(self) -> Dict:
+        return self.dataset[0]
+
+    def __getitem__(self, i) -> Dict:
+        try:
+            item = self.dataset[i]
+            use_multi_task = MultiTaskType.NO_MULTI_TASK
+            for m in self.modalities:
+                if m.use_multi_task != MultiTaskType.NO_MULTI_TASK:
+                    use_multi_task = m.use_multi_task
+                    break
+            if use_multi_task != MultiTaskType.NO_MULTI_TASK:
+                return encode_chat_multitask(item, self.tokenizer, self.modalities)
+            else:
+                return encode_chat(item, self.tokenizer, self.modalities)
+        except Exception as e:
+            new_i = i + 1
+            if new_i >= len(self):
+                new_i = 0
+            logging.error(f"Error encoding chat: {e} index={i} trying index={new_i}")
+            return self.__getitem__(new_i)
+
+
+@dataclass
+class DataCollatorForSupervisedLMMDataset:
+    def __init__(
+        self,
+        tokenizer: transformers.PreTrainedTokenizer,
+        modalities: List[Modality],
+    ):
+        self.tokenizer = tokenizer
+        self.modalities = modalities
+
+        self.use_multi_task = MultiTaskType.NO_MULTI_TASK
+        for modality in self.modalities:
+            if modality.use_multi_task != MultiTaskType.NO_MULTI_TASK:
+                self.use_multi_task = modality.use_multi_task
+                break
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, List]:
+        input_ids = []
+        lmm_labels = []
+        task_labels = []
+        for instance in instances:
+            input_ids.append(instance["input_ids"])
+            if self.use_multi_task == MultiTaskType.NO_MULTI_TASK:
+                lmm_labels.append(instance["labels"])
+            else:
+                lmm_labels.append(instance["labels"][0])
+                inst_task_labels = []
+                for label_id in range(1, len(instance["labels"])):
+                    inst_task_labels.append(instance["labels"][label_id])
+                task_labels.append(inst_task_labels)
+
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        # print("Lmm labels 1 type :", type(lmm_labels))
+        lmm_labels = torch.nn.utils.rnn.pad_sequence(
+            lmm_labels, batch_first=True, padding_value=IGNORE_INDEX
+        )
+        # print("Lmm labels 2 type :", type(lmm_labels))
+
+        input_ids = input_ids[:, : self.tokenizer.model_max_length]
+        lmm_labels = lmm_labels[:, : self.tokenizer.model_max_length]
+        output_labels = [lmm_labels, task_labels]
+        batch = dict(
+            input_ids=input_ids,
+            labels=output_labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+
+        for m in self.modalities:
+            batch[m.name] = [instance[m.name] for instance in instances]
+
+        return batch

src/sonicverse/requirements.txt

@@ -0,0 +1,8 @@
+transformers>=4.34.0
+accelerate>=0.21.0
+scipy>=1.11.3
+bitsandbytes>=0.41.0
+datasets>=2.14.5
+sentencepiece>=0.1.99
+peft>=0.4.0
+deepspeed==0.9.5

src/sonicverse/scripts/audio_setup.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+pip install librosa soundfile

src/sonicverse/scripts/clap_gpt_build_finetune_dataset.py

@@ -0,0 +1,155 @@
+from typing import List
+import argparse
+import json
+import os
+import random
+import openai
+
+from datasets import Dataset, load_dataset
+
+from multi_token.constants import ROLE_ASSISTANT, ROLE_USER
+
+PROMPT = """
+You are helping train a sound assistant that can take audio inputs and output text.
+
+You can hear an audio file with the following metadata tags:
+{captions}
+
+{question}
+
+Include the question and answer.
+"""
+
+QUESTIONS = [
+    "Ask a question about the content of the audio.",
+    "Ask a complex question about the content of the audio.",
+    "Ask a complex question that is relevant to the content of the audio, for example, asking about background knowledge of the things mentioned. Do not ask about uncertain details.",
+    "Ask a complex question that is relevant to the content of the audio, for example, asking about the events referred to in the audio. Do not ask about uncertain details.",
+    "Ask about your thoughts on the audio.",
+    "Ask about what occurs in the audio.",
+    "Ask a question on a topic that related to the audio.",
+    "Ask a question that classifies the audio in some way.",
+    "Ask a question that can only be answered by listening to the audio.",
+]
+
+
+OPENAI_TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "create_chat",
+            "description": "Create a training example",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "question": {
+                        "type": "string",
+                        "description": "The question, must be provided",
+                    },
+                    "answer": {
+                        "type": "string",
+                        "description": "The answer to the question, must be provided",
+                    },
+                },
+                "required": ["question", "answer"],
+            },
+        },
+    }
+]
+
+
+def _build_convo(row) -> List:
+    client = openai.Client()
+
+    captions = [row["metadataTags"]]
+    paths = [row["url"]]
+
+    captions_text = "\n".join([f"{cap}" for i, cap in enumerate(captions)])
+    prompt = PROMPT.format(
+        captions=captions_text, question=random.choice(QUESTIONS)
+    ).strip()
+
+    completion = client.chat.completions.create(
+        model="gpt-3.5-turbo-1106",
+        messages=[{"role": "system", "content": prompt}],
+        tools=OPENAI_TOOLS,
+        tool_choice={"type": "function", "function": {"name": "create_chat"}},
+    )
+    resp = json.loads(completion.choices[0].message.tool_calls[0].function.arguments)
+    if "answer" not in resp:
+        print(resp)
+    q = resp["question"]
+    a = resp["answer"]
+
+    if random.choice([True, False]):
+        q = "<sound>" * len(captions) + " " + q
+    else:
+        q = q + " " + "<sound>" * len(captions)
+
+    example = {
+        "sounds": paths,
+        "messages": [
+            {
+                "role": ROLE_USER,
+                "content": q,
+            },
+            {
+                "role": ROLE_ASSISTANT,
+                "content": a,
+            },
+        ],
+    }
+    return example
+
+
+def main(args):
+    data = load_dataset("Chr0my/Epidemic_sounds", split="train")
+    data_idxs = list(range(len(data)))
+
+    os.makedirs(args.cache_folder, exist_ok=True)
+
+    def gen(seeds):
+        r = random.Random(seeds[0] + 3)
+        cache = open(
+            os.path.join(args.cache_folder, f"gpt-cache.{seeds[0]}.jsonl"), "a"
+        )
+        i = 0
+        while i < len(seeds):
+            selected_idxs = r.sample(data_idxs, k=1)[0]
+            selected_example = data[selected_idxs]
+            try:
+                example = _build_convo(selected_example)
+                cache.write(json.dumps(example) + "\n")
+                yield example
+                i += 1
+            except Exception as e:
+                print(e)
+                continue
+        cache.close()
+
+    ds = Dataset.from_generator(
+        gen,
+        num_proc=args.num_proc,
+        gen_kwargs={"seeds": list(range(args.num_examples))},
+    )
+    ds.save_to_disk(args.output_folder)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        "--output_folder",
+        type=str,
+        default="/data/clap-gpt-finetune",
+    )
+    parser.add_argument(
+        "-c",
+        "--cache_folder",
+        type=str,
+        default="/data/clap-gpt-finetune-cache",
+    )
+    parser.add_argument("-n", "--num_examples", type=int, default=100_000)
+    parser.add_argument("-p", "--num_proc", type=int, default=10)
+    args = parser.parse_args()
+    main(args)

src/sonicverse/scripts/clap_gpt_build_pretrain_dataset.py

@@ -0,0 +1,142 @@
+from typing import List
+import argparse
+import json
+import os
+import random
+import openai
+
+from datasets import Dataset, load_dataset
+
+from multi_token.constants import ROLE_ASSISTANT, ROLE_USER
+
+PROMPT = """
+You are helping write captions for audio clips.
+
+Here are the tags for the audio clip you are captioning:
+{captions}
+
+Write a brief caption for the audio clip.
+"""
+
+PRETRAIN_PHRASES = [
+    "What is happening in <sound>?",
+    "Describe the sound. <sound>",
+    "<sound> Provide a description of the audio.",
+    "Can you interpret <sound>?",
+    "Please explain what's happening in <sound>",
+    "What does <sound> represent?",
+    "Could you describe <sound> for me?",
+    "What's the content of <sound>?",
+    "Can you depict <sound>?",
+    "What is <sound>?",
+    "In the audo clip, <sound>, what is happening?",
+    "Provide a description of the sound. <sound>",
+    "Provide a caption for the sound. <sound>",
+]
+
+OPENAI_TOOLS = [
+    {
+        "type": "function",
+        "function": {
+            "name": "write_caption",
+            "description": "Write a caption for an audio clip",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "caption": {
+                        "type": "string",
+                    },
+                },
+                "required": ["caption"],
+            },
+        },
+    }
+]
+
+
+def _build_convo(row) -> List:
+    client = openai.Client()
+
+    captions = [row["metadataTags"]]
+    sounds = [row["url"]]
+
+    captions_text = "\n".join([f'Tags: "{cap}"' for i, cap in enumerate(captions)])
+    prompt = PROMPT.format(captions=captions_text).strip()
+
+    completion = client.chat.completions.create(
+        model="gpt-3.5-turbo-1106",
+        messages=[{"role": "system", "content": prompt}],
+        tools=OPENAI_TOOLS,
+        tool_choice={"type": "function", "function": {"name": "write_caption"}},
+    )
+    resp = json.loads(completion.choices[0].message.tool_calls[0].function.arguments)
+    caption = resp["caption"]
+
+    q = random.choice(PRETRAIN_PHRASES)
+
+    example = {
+        "sounds": sounds,
+        "messages": [
+            {
+                "role": ROLE_USER,
+                "content": q,
+            },
+            {
+                "role": ROLE_ASSISTANT,
+                "content": caption,
+            },
+        ],
+    }
+    return example
+
+
+def main(args):
+    data = load_dataset("Chr0my/Epidemic_sounds", split="train")
+
+    os.makedirs(args.cache_folder, exist_ok=True)
+
+    def gen(seeds):
+        cache = open(
+            os.path.join(args.cache_folder, f"gpt-cache.{seeds[0]}.jsonl"), "a"
+        )
+        for s in seeds:
+            selected_row = data[s]
+            try:
+                example = _build_convo(selected_row)
+                cache.write(json.dumps(example) + "\n")
+                yield example
+            except Exception as e:
+                print(e)
+                continue
+
+        cache.close()
+
+    idxs = list(range(len(data)))
+    random.shuffle(idxs)
+
+    ds = Dataset.from_generator(
+        gen,
+        num_proc=args.num_proc,
+        gen_kwargs={"seeds": idxs},
+    )
+    ds.save_to_disk(args.output_folder)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-o",
+        "--output_folder",
+        type=str,
+        default="/data/clap-gpt-pretrain",
+    )
+    parser.add_argument(
+        "-c",
+        "--cache_folder",
+        type=str,
+        default="/data/clap-gpt-pretrain-cache",
+    )
+    parser.add_argument("-n", "--num_examples", type=int, default=500_000)
+    parser.add_argument("-p", "--num_proc", type=int, default=10)
+    args = parser.parse_args()
+    main(args)

src/sonicverse/scripts/document_build_finetune_dataset.py

@@ -0,0 +1,162 @@
+from typing import List
+import argparse
+import re
+import glob
+import json
+
+from datasets import load_dataset
+from datasets import Dataset
+
+from multi_token.constants import ROLE_ASSISTANT, ROLE_USER
+from multi_token.modalities.document_gte import (
+    split_text_into_documents,
+)
+
+TEMP_TOKEN = "<<<TEMP-TOKEN>>>"
+
+# regex, doc, prompt
+LONG_ALPACA_REGEXES = [
+    (
+        r"Below is a paper. Memorize the paper and answer my question after the paper.\n The paper begins. \n ([\s\S]+) \n Now the paper ends. \n([\s\S]+)",
+        lambda m: m.group(1),
+        lambda m: f"Read the paper {TEMP_TOKEN}. {m.group(2)}",
+    ),
+    (
+        r"Below is a paper. Memorize the material and answer my question after the paper.\n([\s\S]+)\n Now the material ends. ([\s\S]+)",
+        lambda m: m.group(1),
+        lambda m: f"Read the paper {TEMP_TOKEN}. {m.group(2)}",
+    ),
+    (
+        r"There are two papers. Memorize them and answer my question after the paper.\n The first paper begins. \n ([\s\S]+) Now the second paper ends.([\s\S]+)",
+        lambda m: m.group(1),
+        lambda m: f"Read the papers {TEMP_TOKEN}. {m.group(2)}",
+    ),
+    (
+        r"Below is some paragraphs in the book, ([\s\S]+?). Memorize the content and answer my question after the book.\n([\s\S]+) \n Now the material ends.([\s\S]+)",
+        lambda m: m.group(2),
+        lambda m: f"Read the book {m.group(1)} {TEMP_TOKEN}. {m.group(3)}",
+    ),
+]
+
+# regex, doc, prompt, answer
+LONG_DATA_REGEXES = [
+    (
+        r"Write a high-quality answer for the given question using only the provided search results \(some of which might be irrelevant\).([\s\S]+)Question: ([\s\S]+)Answer: ([\s\S]+)\nLong Answer: ([\s\S]+)",
+        lambda m: m.group(1).strip(),
+        lambda m: f"Write a high-quality answer for the given question using only the provided search results {TEMP_TOKEN}. {m.group(2).strip()}",
+        lambda m: m.group(4).strip(),
+    ),
+    (
+        r"([\s\S]+)\nQ: ([\s\S]+)\nA: ([\s\S]+)",
+        lambda m: m.group(1).strip(),
+        lambda m: f"Read the following book {TEMP_TOKEN}. {m.group(2).strip()}",
+        lambda m: m.group(3).strip(),
+    ),
+]
+
+
+def _write_long_alpaca_convo(row, max_document_chunks) -> List:
+    doc_text = None
+    prompt = None
+    for regex, get_doc, get_prompt in LONG_ALPACA_REGEXES:
+        match = re.match(regex, row["instruction"])
+        if match:
+            doc_text = get_doc(match)
+            prompt = get_prompt(match).replace("Question: ", "")
+            break
+
+    if doc_text is None and row["input"]:
+        doc_text = row["input"]
+        prompt = row["instruction"] + f" {TEMP_TOKEN}"
+
+    if doc_text is None:
+        raise ValueError("No document found")
+
+    docs = split_text_into_documents(doc_text)
+    if len(docs) > max_document_chunks:
+        raise ValueError("Document too long")
+    example = {
+        "id": "longalpaca-" + str(hash(row["instruction"])),
+        "documents": docs,
+    }
+    example["messages"] = [
+        {
+            "role": ROLE_USER,
+            "content": prompt.replace(TEMP_TOKEN, "<document>" * len(docs)),
+        },
+        {
+            "role": ROLE_ASSISTANT,
+            "content": row["output"].replace("Answer: ", ""),
+        },
+    ]
+    return example
+
+
+def _write_long_data_collections_convo(row, max_document_chunks) -> List:
+    doc_text = None
+    prompt = None
+    answer = None
+    for regex, get_doc, get_prompt, get_answer in LONG_DATA_REGEXES:
+        match = re.match(regex, row["text"])
+        if match:
+            doc_text = get_doc(match)
+            prompt = get_prompt(match)
+            answer = get_answer(match).replace(" .", ".")
+            break
+
+    if not doc_text or not prompt or not answer:
+        raise ValueError("No document found")
+
+    docs = split_text_into_documents(doc_text)
+    if len(docs) > max_document_chunks:
+        raise ValueError("Document too long")
+    example = {
+        "id": "longdatacollection-" + str(hash(row["text"])),
+        "documents": docs,
+    }
+    example["messages"] = [
+        {
+            "role": ROLE_USER,
+            "content": prompt.replace(TEMP_TOKEN, "<document>" * len(docs)),
+        },
+        {
+            "role": ROLE_ASSISTANT,
+            "content": answer,
+        },
+    ]
+    return example
+
+
+def main(args):
+    long_alpaca = load_dataset(args.long_alpaca_path, "train")["train"]
+
+    def gen():
+        for row in long_alpaca:
+            try:
+                yield _write_long_alpaca_convo(row, args.max_document_chunks)
+            except ValueError:
+                continue
+        for long_collection_fn in glob.iglob(args.long_collections_glob):
+            with open(long_collection_fn) as f:
+                for line in f:
+                    row = json.loads(line)
+                    try:
+                        yield _write_long_data_collections_convo(
+                            row, args.max_document_chunks
+                        )
+                    except ValueError:
+                        continue
+
+    ds = Dataset.from_generator(gen)
+    ds = ds.shuffle(seed=42)
+    ds.save_to_disk(args.output_folder)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--long_alpaca_path", type=str, default="Yukang/LongAlpaca-12k")
+    parser.add_argument("--long_collections_glob", type=str)
+    parser.add_argument("-o", "--output_folder", type=str)
+    parser.add_argument("-c", "--max_document_chunks", type=int, default=256)
+    args = parser.parse_args()
+    main(args)

src/sonicverse/scripts/document_build_pretrain_dataset.py

@@ -0,0 +1,89 @@
+from typing import List
+import random
+import argparse
+
+from datasets import load_dataset
+from datasets import Dataset
+
+from multi_token.constants import ROLE_ASSISTANT, ROLE_USER
+from multi_token.modalities.document_gte import (
+    split_text_into_documents,
+)
+
+TEMP_TOKEN = "<<<TEMP-TOKEN>>>"
+
+PRETRAIN_PHRASES = [
+    f"Repeat the content of the document {TEMP_TOKEN}",
+    f"Transcribe {TEMP_TOKEN}",
+    f"Provide a verbatim transcription of {TEMP_TOKEN}",
+    f"Write down exactly what is in {TEMP_TOKEN}",
+    f"Copy the text from {TEMP_TOKEN}",
+    f"Duplicate the content of {TEMP_TOKEN}",
+    f"Reproduce the text in {TEMP_TOKEN}",
+    f"Render the exact text from {TEMP_TOKEN}",
+    f"Echo the content of {TEMP_TOKEN}",
+    f"Mirror the text in {TEMP_TOKEN}",
+    f"Reflect the content of {TEMP_TOKEN}",
+    f"Transcribe the exact words from {TEMP_TOKEN}",
+    f"Write out the exact content of {TEMP_TOKEN}",
+    f"Provide a direct transcription of {TEMP_TOKEN}",
+    f"Give a word-for-word account of {TEMP_TOKEN}",
+    f"Reiterate the exact text of {TEMP_TOKEN}",
+    f"Replicate the content of {TEMP_TOKEN}",
+    f"Reprint the text from {TEMP_TOKEN}",
+    f"Rewrite the exact words from {TEMP_TOKEN}",
+]
+
+
+def _write_convo(row, max_document_chunks) -> List:
+    docs = split_text_into_documents(row["text"])
+    if len(docs) > max_document_chunks:
+        raise ValueError("Document too long")
+    example = {
+        "id": str(row["title"]),
+        "documents": docs,
+    }
+    phrase = random.choice(PRETRAIN_PHRASES)
+    example["messages"] = [
+        {
+            "role": ROLE_USER,
+            "content": phrase.replace(TEMP_TOKEN, "<document>" * len(docs)),
+        },
+        {
+            "role": ROLE_ASSISTANT,
+            "content": row["text"],
+        },
+    ]
+    return example
+
+
+def main(args):
+    wiki_data = load_dataset("graelo/wikipedia", "20230601.en")["train"]
+
+    idxs = list(range(len(wiki_data)))
+    random.shuffle(idxs)
+
+    def gen():
+        i = 0
+        for idx in idxs:
+            row = wiki_data[idx]
+            try:
+                yield _write_convo(row, args.max_document_chunks)
+            except ValueError:
+                pass
+            else:
+                i += 1
+                if i >= args.max_examples:
+                    break
+
+    ds = Dataset.from_generator(gen)
+    ds.save_to_disk(args.output_folder)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--output_folder", type=str)
+    parser.add_argument("-n", "--max_examples", type=int, default=1_000_000)
+    parser.add_argument("-c", "--max_document_chunks", type=int, default=4)
+    args = parser.parse_args()
+    main(args)

src/sonicverse/scripts/document_setup.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+
+pip install nltk
+
+python -c "import nltk; nltk.download('punkt')"

src/sonicverse/scripts/evaluate_model.py

@@ -0,0 +1,112 @@
+from dataclasses import dataclass, field
+import logging
+
+from flask import Flask, request, jsonify
+import transformers
+import torch
+
+from datasets import load_from_disk
+
+from multi_token.model_utils import MultiTaskType
+from multi_token.training import (
+    ModelArguments,
+)
+from multi_token.inference import load_trained_lora_model
+from multi_token.data_tools import encode_chat
+
+import evaluate
+
+import random
+
+PRETRAIN_PHRASES = [
+    "What is happening in the given music <sound>?",
+    "Describe the sound. <sound>",
+    "Describe the music. <sound>",
+    "<sound> Provide a description of the music.",
+    "<sound> Provide a description of the sound.",
+    "Can you interpret <sound>?",
+    "Please explain what's happening in <sound>",
+    "What does <sound> represent?",
+    "Could you describe <sound> for me?",
+    "What's the content of <sound>?",
+    "Can you depict <sound>?",
+    "What is <sound>?",
+    "In the music clip, <sound>, what is happening?",
+    "Provide a description of the music. <sound>",
+    "Provide a description of the sound. <sound>",
+    "Provide a caption for the sound. <sound>",
+    "Provide a caption for the music. <sound>",
+]
+
+
+@dataclass
+class ServeArguments(ModelArguments):
+    port: int = field(default=8080)
+    host: str = field(default="0.0.0.0")
+    load_bits: int = field(default=16)
+    max_new_tokens: int = field(default=128)
+    temperature: float = field(default=0.01)
+
+
+def generate(input_json):
+    encoded_dict = encode_chat(input_json, tokenizer, model.modalities)
+
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
+            max_new_tokens=serve_args.max_new_tokens,
+            use_cache=True,
+            do_sample=True,
+            temperature=serve_args.temperature,
+            modality_inputs={
+                m.name: [encoded_dict[m.name]] for m in model.modalities
+            },
+        )
+
+    outputs = tokenizer.decode(
+        output_ids[0, encoded_dict["input_ids"].shape[0] :],
+        skip_special_tokens=True,
+    ).strip()
+
+    return {"output": outputs}
+
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+
+    parser = transformers.HfArgumentParser((ServeArguments,))
+
+    serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+
+    dataset_path = "/data/musicbench_multitoken_official_split/val"
+
+    ds = load_from_disk(dataset_path)
+
+    model, tokenizer = load_trained_lora_model(
+        model_name_or_path=serve_args.model_name_or_path,
+        model_lora_path=serve_args.model_lora_path,
+        load_bits=serve_args.load_bits,
+        use_multi_task=MultiTaskType(serve_args.use_multi_task),
+        tasks_config=serve_args.tasks_config
+    )
+
+    predictions = []
+    references = []
+    content_phrase = random.choice(PRETRAIN_PHRASES)
+    for data_point_id in range(100):
+        data_point = ds[data_point_id]
+        # print("datapoint", data_point)
+        input_json={"messages": [{"role": "user", "content": content_phrase}], "sounds": data_point["sounds"]}
+        output_json = generate(input_json)
+
+        print("Prediction ",output_json["output"])
+        print("Reference ", data_point["messages"][1]["content"])
+        print()
+        print()
+        predictions.append(output_json["output"])
+        references.append(data_point["messages"][1]["content"])
+
+    sacrebleu = evaluate.load("sacrebleu")
+    sacrebleu_results=sacrebleu.compute(predictions=predictions, references=references)
+
+    print(sacrebleu_results["score"])

src/sonicverse/scripts/evaluate_model_latest.py

@@ -0,0 +1,127 @@
+
+from dataclasses import dataclass, field
+import logging
+from flask import Flask, request, jsonify
+import transformers
+import torch
+from datasets import load_from_disk
+from multi_token.model_utils import MultiTaskType
+from multi_token.training import ModelArguments
+from multi_token.inference import load_trained_lora_model
+from multi_token.data_tools import encode_chat
+import evaluate
+import random
+import bert_score
+from tqdm import tqdm
+
+PRETRAIN_PHRASES_OLD = [
+    "Describe the audio in detail"
+]
+
+PRETRAIN_PHRASES = [
+    "What is happening in the given music <sound>?",
+    "Describe the sound. <sound>",
+    "Describe the music. <sound>",
+    "<sound> Provide a description of the music.",
+    "<sound> Provide a description of the sound.",
+    "Can you interpret <sound>?",
+    "Please explain what's happening in <sound>",
+    "What does <sound> represent?",
+    "Could you describe <sound> for me?",
+    "What's the content of <sound>?",
+    "Can you depict <sound>?",
+    "What is <sound>?",
+    "In the music clip, <sound>, what is happening?",
+    "Provide a description of the music. <sound>",
+    "Provide a description of the sound. <sound>",
+    "Provide a caption for the sound. <sound>",
+    "Provide a caption for the music. <sound>",
+]
+
+random.seed(1234)
+
+@dataclass
+class ServeArguments(ModelArguments):
+    port: int = field(default=8080)
+    host: str = field(default="0.0.0.0")
+    load_bits: int = field(default=16)
+    max_new_tokens: int = field(default=128)
+    temperature: float = field(default=0.01)
+
+def generate(input_json):
+    encoded_dict = encode_chat(input_json, tokenizer, model.modalities)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
+            max_new_tokens=serve_args.max_new_tokens,
+            use_cache=True,
+            do_sample=True,
+            temperature=serve_args.temperature,
+            modality_inputs={
+                m.name: [encoded_dict[m.name]] for m in model.modalities
+            },
+        )
+    outputs = tokenizer.decode(
+        output_ids[0, encoded_dict["input_ids"].shape[0]:],
+        skip_special_tokens=True,
+    ).strip()
+    return {"output": outputs}
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    parser = transformers.HfArgumentParser((ServeArguments,))
+    serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    dataset_path = "/data/musicbench_multitoken_official_split/val"
+    ds = load_from_disk(dataset_path)
+    shuffled_ds = ds.shuffle(seed=1234)
+    model, tokenizer = load_trained_lora_model(
+        model_name_or_path=serve_args.model_name_or_path,
+        model_lora_path=serve_args.model_lora_path,
+        load_bits=serve_args.load_bits,
+        use_multi_task=MultiTaskType(serve_args.use_multi_task),
+        tasks_config=serve_args.tasks_config
+    )
+
+    predictions = []
+    references = []
+    content_phrase = random.choice(PRETRAIN_PHRASES)
+#    for data_point_id in range(len(ds)):
+    for data_point_id in tqdm(range(10)):
+        data_point = shuffled_ds[data_point_id]
+        input_json = {"messages": [{"role": "user", "content": content_phrase}], "sounds": data_point["sounds"]}
+        output_json = generate(input_json)
+        print("Prediction ", output_json["output"])
+        print("Reference ", data_point["messages"][1]["content"])
+        print()
+        print()
+        predictions.append(output_json["output"])
+        references.append(data_point["messages"][1]["content"])
+
+    # Load evaluation metrics
+    bleu = evaluate.load("bleu")
+    meteor = evaluate.load("meteor")
+    rouge = evaluate.load("rouge")
+    
+    # Compute BLEU scores
+    bleu_results = bleu.compute(predictions=predictions, references=references, max_order=4)
+    print(bleu_results)
+    #bleu_score = sum(bleu_results[f"bleu{i}"] for i in range(1, 5)) / 4
+
+    # Compute METEOR score
+    meteor_results = meteor.compute(predictions=predictions, references=references)
+    meteor_score = meteor_results["meteor"]
+
+    # Compute ROUGE-L score
+    rouge_results = rouge.compute(predictions=predictions, references=references, rouge_types=["rougeL"])
+    rouge_l_score = rouge_results["rougeL"].mid.fmeasure
+    #print(rouge_results)
+
+    # Compute BERT-Score
+    P, R, F1 = bert_score.score(predictions, references, lang="en", rescale_with_baseline=True)
+    bert_score_f1 = F1.mean().item()
+
+    # Print results
+    #print(f"BLEU Score: {bleu_score}")
+    print(f"METEOR Score: {meteor_score}")
+    #print(f"ROUGE-L Score: {rouge_l_score}")
+    print(f"BERT-Score F1: {bert_score_f1}")

src/sonicverse/scripts/evaluate_model_mullama.py

@@ -0,0 +1,168 @@
+
+from dataclasses import dataclass, field
+import logging
+from flask import Flask, request, jsonify
+import transformers
+import torch
+from datasets import load_from_disk
+from multi_token.model_utils import MultiTaskType
+from multi_token.training import ModelArguments
+from multi_token.inference import load_trained_lora_model
+from multi_token.data_tools import encode_chat
+import evaluate
+import random
+import bert_score
+from tqdm import tqdm
+
+from rouge_score import rouge_scorer
+from nltk.translate.bleu_score import sentence_bleu
+from nltk.translate.meteor_score import meteor_score as meteor_scorer
+from nltk.tokenize import wordpunct_tokenize
+import json
+from bert_score import score
+from tqdm.auto import tqdm
+
+import yaml
+
+scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
+
+
+PRETRAIN_PHRASES_OLD = [
+    "Describe the audio in detail"
+]
+
+PRETRAIN_PHRASES = [
+    "What is happening in the given music <sound>?",
+    "Describe the sound. <sound>",
+    "Describe the music. <sound>",
+    "<sound> Provide a description of the music.",
+    "<sound> Provide a description of the sound.",
+    "Can you interpret <sound>?",
+    "Please explain what's happening in <sound>",
+    "What does <sound> represent?",
+    "Could you describe <sound> for me?",
+    "What's the content of <sound>?",
+    "Can you depict <sound>?",
+    "What is <sound>?",
+    "In the music clip, <sound>, what is happening?",
+    "Provide a description of the music. <sound>",
+    "Provide a description of the sound. <sound>",
+    "Provide a caption for the sound. <sound>",
+    "Provide a caption for the music. <sound>",
+]
+
+random.seed(1234)
+
+@dataclass
+class ServeArguments(ModelArguments):
+    port: int = field(default=8080)
+    host: str = field(default="0.0.0.0")
+    load_bits: int = field(default=16)
+    max_new_tokens: int = field(default=128)
+    temperature: float = field(default=0.01)
+
+def generate(input_json):
+    encoded_dict = encode_chat(input_json, tokenizer, model.modalities)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
+            max_new_tokens=serve_args.max_new_tokens,
+            use_cache=True,
+            do_sample=True,
+            temperature=serve_args.temperature,
+            modality_inputs={
+                m.name: [encoded_dict[m.name]] for m in model.modalities
+            },
+        )
+    outputs = tokenizer.decode(
+        output_ids[0, encoded_dict["input_ids"].shape[0]:],
+        skip_special_tokens=True,
+    ).strip()
+    return {"output": outputs}
+
+def evaluate(candidates, mult_reference):
+    rouge_score, bleu_score, bleu4_score, meteor_score = 0, 0, 0, 0
+    for ref, cand in tqdm(zip(mult_reference, candidates), total=len(mult_reference)):
+        rouge_score += scorer.score(ref, cand)['rougeL'].recall
+        cand_split = wordpunct_tokenize(cand)
+        ref_split = wordpunct_tokenize(ref)
+        bleu4_score += sentence_bleu([ref], cand, weights=(0.0, 0.0, 0.0, 1.0))
+        bleu_score += sentence_bleu([ref], cand)
+        meteor_score += meteor_scorer([ref_split], cand_split)
+    rouge_score, bleu_score, bleu4_score, meteor_score = rouge_score / (len(candidates)), bleu_score / (len(candidates)), bleu4_score / (len(candidates)), meteor_score / (len(candidates))
+    P, R, F1 = score(candidates, mult_reference, lang="en", verbose=True)
+    bert_score = R.mean().item()
+    #print(f"Model: {model_name}")
+    print(f"BLEU Score: {bleu_score}")
+    print(f"BLEU-4 Score: {bleu4_score}")
+    print(f"METEOR Score: {meteor_score}")
+    print(f"ROUGE Score: {rouge_score}")
+    print(f"BERT Score: {bert_score}")
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    parser = transformers.HfArgumentParser((ServeArguments,))
+    serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    dataset_path = "/data/musicbench_multitoken_official_split/val"
+    ds = load_from_disk(dataset_path)
+    shuffled_ds = ds.shuffle(seed=1234)
+    model, tokenizer = load_trained_lora_model(
+        model_name_or_path=serve_args.model_name_or_path,
+        model_lora_path=serve_args.model_lora_path,
+        load_bits=serve_args.load_bits,
+        use_multi_task=MultiTaskType(serve_args.use_multi_task),
+        tasks_config=serve_args.tasks_config
+    )
+
+    predictions = []
+    references = []
+    content_phrase = random.choice(PRETRAIN_PHRASES)
+#    for data_point_id in range(len(ds)):
+    print("len(ds)", len(ds))
+    for data_point_id in tqdm(range(100)):
+#    for data_point_id in tqdm(range(6831)):
+        data_point = shuffled_ds[data_point_id]
+        input_json = {"messages": [{"role": "user", "content": content_phrase}], "sounds": data_point["sounds"]}
+        output_json = generate(input_json)
+#        print("Prediction ", output_json["output"])
+#        print("Reference ", data_point["messages"][1]["content"])
+#        print()
+#        print()
+        predictions.append(output_json["output"])
+        references.append(data_point["messages"][1]["content"])
+
+    pairs = {"predictions": predictions, "references": references}
+
+    evaluate(predictions, references)
+
+#    with open('/experiments/captioning/mert_tasks_separate_backbone_train_001_ft/checkpoint_1985_test/val_2.yaml', 'w') as file:
+#        yaml.dump(pairs, file, default_flow_style=False)
+
+    # Load evaluation metrics
+    # bleu = evaluate.load("bleu")
+    # meteor = evaluate.load("meteor")
+    # rouge = evaluate.load("rouge")
+    
+    # Compute BLEU scores
+    # bleu_results = bleu.compute(predictions=predictions, references=references, max_order=4)
+    # print(bleu_results)
+    #bleu_score = sum(bleu_results[f"bleu{i}"] for i in range(1, 5)) / 4
+
+    # Compute METEOR score
+    # meteor_results = meteor.compute(predictions=predictions, references=references)
+    # meteor_score = meteor_results["meteor"]
+
+    # Compute ROUGE-L score
+    # rouge_results = rouge.compute(predictions=predictions, references=references, rouge_types=["rougeL"])
+#    rouge_l_score = rouge_results["rougeL"].mid.fmeasure
+    # print(rouge_results)
+
+    # Compute BERT-Score
+    # P, R, F1 = bert_score.score(predictions, references, lang="en", rescale_with_baseline=True)
+    # bert_score_f1 = F1.mean().item()
+
+    # Print results
+    #print(f"BLEU Score: {bleu_score}")
+    # print(f"METEOR Score: {meteor_score}")
+#    print(f"ROUGE-L Score: {rouge_l_score}")
+    # print(f"BERT-Score F1: {bert_score_f1}")

src/sonicverse/scripts/evaluate_model_mullama_musiccaps.py

@@ -0,0 +1,143 @@
+
+from dataclasses import dataclass, field
+import logging
+from flask import Flask, request, jsonify
+import transformers
+import torch
+from datasets import load_from_disk
+from multi_token.model_utils import MultiTaskType
+from multi_token.training import ModelArguments
+from multi_token.inference import load_trained_lora_model
+from multi_token.data_tools import encode_chat
+import evaluate
+import random
+import bert_score
+from tqdm import tqdm
+
+from rouge_score import rouge_scorer
+from nltk.translate.bleu_score import sentence_bleu
+from nltk.translate.meteor_score import meteor_score as meteor_scorer
+from nltk.tokenize import wordpunct_tokenize
+import json
+from bert_score import score
+from tqdm.auto import tqdm
+
+import yaml
+
+scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
+
+
+PRETRAIN_PHRASES_OLD = [
+    "Describe the audio in detail"
+]
+
+PRETRAIN_PHRASES = [
+    "What is happening in the given music <sound>?",
+    "Describe the sound. <sound>",
+    "Describe the music. <sound>",
+    "<sound> Provide a description of the music.",
+    "<sound> Provide a description of the sound.",
+    "Can you interpret <sound>?",
+    "Please explain what's happening in <sound>",
+    "What does <sound> represent?",
+    "Could you describe <sound> for me?",
+    "What's the content of <sound>?",
+    "Can you depict <sound>?",
+    "What is <sound>?",
+    "In the music clip, <sound>, what is happening?",
+    "Provide a description of the music. <sound>",
+    "Provide a description of the sound. <sound>",
+    "Provide a caption for the sound. <sound>",
+    "Provide a caption for the music. <sound>",
+]
+
+random.seed(1234)
+
+@dataclass
+class ServeArguments(ModelArguments):
+    port: int = field(default=8080)
+    host: str = field(default="0.0.0.0")
+    load_bits: int = field(default=16)
+    max_new_tokens: int = field(default=128)
+    temperature: float = field(default=0.01)
+
+def generate(input_json):
+    encoded_dict = encode_chat(input_json, tokenizer, model.modalities)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
+            max_new_tokens=serve_args.max_new_tokens,
+            use_cache=True,
+            do_sample=True,
+            temperature=serve_args.temperature,
+            modality_inputs={
+                m.name: [encoded_dict[m.name]] for m in model.modalities
+            },
+        )
+    outputs = tokenizer.decode(
+        output_ids[0, encoded_dict["input_ids"].shape[0]:],
+        skip_special_tokens=True,
+    ).strip()
+    return {"output": outputs}
+
+def evaluate(candidates, mult_reference):
+    rouge_score, bleu_score, bleu4_score, meteor_score = 0, 0, 0, 0
+    for ref, cand in tqdm(zip(mult_reference, candidates), total=len(mult_reference)):
+        rouge_score += scorer.score(ref, cand)['rougeL'].recall
+        cand_split = wordpunct_tokenize(cand)
+        ref_split = wordpunct_tokenize(ref)
+        bleu4_score += sentence_bleu([ref], cand, weights=(0.0, 0.0, 0.0, 1.0))
+        bleu_score += sentence_bleu([ref], cand)
+        meteor_score += meteor_scorer([ref_split], cand_split)
+    rouge_score, bleu_score, bleu4_score, meteor_score = rouge_score / (len(candidates)), bleu_score / (len(candidates)), bleu4_score / (len(candidates)), meteor_score / (len(candidates))
+    P, R, F1 = score(candidates, mult_reference, lang="en", verbose=True)
+    bert_score = R.mean().item()
+    #print(f"Model: {model_name}")
+    print(f"BLEU Score: {bleu_score}")
+    print(f"BLEU-4 Score: {bleu4_score}")
+    print(f"METEOR Score: {meteor_score}")
+    print(f"ROUGE Score: {rouge_score}")
+    print(f"BERT Score: {bert_score}")
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    parser = transformers.HfArgumentParser((ServeArguments,))
+    serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    # dataset_path = "/data/musiccaps/musiccaps_val"
+    dataset_path = "/data/musicbench_multitoken_official_split/val/"
+    ds = load_from_disk(dataset_path)
+    shuffled_ds = ds.shuffle(seed=1234)
+    model, tokenizer = load_trained_lora_model(
+        model_name_or_path=serve_args.model_name_or_path,
+        model_lora_path=serve_args.model_lora_path,
+        load_bits=serve_args.load_bits,
+        use_multi_task=MultiTaskType(serve_args.use_multi_task),
+        tasks_config=serve_args.tasks_config
+    )
+
+    predictions = []
+    references = []
+    content_phrase = random.choice(PRETRAIN_PHRASES)
+#    for data_point_id in range(len(ds)):
+    print("len(ds)", len(ds))
+    #for data_point in tqdm(ds):
+    for data_point_id in tqdm(range(100)):
+        #print("DATA POINT ", data_point)
+        data_point = ds[data_point_id]
+        print("DATA POINT ", data_point)
+        input_json = {"messages": [{"role": "user", "content": content_phrase}], "sounds": data_point["sounds"]}
+        output_json = generate(input_json)
+        #print("Prediction ", output_json["output"])
+        #print("Reference ", data_point["caption"])
+        #print()
+        #print()
+        predictions.append(output_json["output"])
+        references.append(data_point["messages"][1]["content"])
+
+    pairs = {"predictions": predictions, "references": references}
+
+    evaluate(predictions, references)
+
+    with open('test/musicbench_eval.yaml', 'w') as file:
+        yaml.dump(pairs, file, default_flow_style=False)
+

src/sonicverse/scripts/evaluate_model_mullama_musiccaps_fixed_prompt.py

@@ -0,0 +1,138 @@
+from dataclasses import dataclass, field
+import logging
+from flask import Flask, request, jsonify
+import transformers
+import torch
+from datasets import load_from_disk
+from multi_token.model_utils import MultiTaskType
+from multi_token.training import ModelArguments
+from multi_token.inference import load_trained_lora_model
+from multi_token.data_tools import encode_chat
+import evaluate
+import random
+import bert_score
+from tqdm import tqdm
+
+from rouge_score import rouge_scorer
+from nltk.translate.bleu_score import sentence_bleu
+from nltk.translate.meteor_score import meteor_score as meteor_scorer
+from nltk.tokenize import wordpunct_tokenize
+import json
+from bert_score import score
+from tqdm.auto import tqdm
+
+import yaml
+
+scorer = rouge_scorer.RougeScorer(['rougeL'], use_stemmer=True)
+
+
+PRETRAIN_PHRASES_OLD = [
+    "Describe the audio in detail"
+]
+
+PRETRAIN_PHRASES = [
+    # "What is happening in the given music <sound>?",
+    # "Describe the sound. <sound>",
+    # "Describe the music. <sound>",
+    # "<sound> Provide a description of the music.",
+    # "<sound> Provide a description of the sound.",
+    # "Can you interpret <sound>?",
+    # "Please explain what's happening in <sound>",
+    # "What does <sound> represent?",
+    # "Could you describe <sound> for me?",
+    # "What's the content of <sound>?",
+    # "Can you depict <sound>?",
+    # "What is <sound>?",
+    # "In the music clip, <sound>, what is happening?",
+    # "Provide a description of the music. <sound>",
+    # "Provide a description of the sound. <sound>",
+    # "Provide a caption for the sound. <sound>",
+    "Provide a caption for the music. <sound>",
+]
+
+random.seed(1234)
+
+@dataclass
+class ServeArguments(ModelArguments):
+    port: int = field(default=8080)
+    host: str = field(default="0.0.0.0")
+    load_bits: int = field(default=16)
+    max_new_tokens: int = field(default=128)
+    temperature: float = field(default=0.01)
+
+def generate(input_json):
+    encoded_dict = encode_chat(input_json, tokenizer, model.modalities)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
+            max_new_tokens=serve_args.max_new_tokens,
+            use_cache=True,
+            do_sample=True,
+            temperature=serve_args.temperature,
+            modality_inputs={
+                m.name: [encoded_dict[m.name]] for m in model.modalities
+            },
+        )
+    outputs = tokenizer.decode(
+        output_ids[0, encoded_dict["input_ids"].shape[0]:],
+        skip_special_tokens=True,
+    ).strip()
+    return {"output": outputs}
+
+def evaluate(candidates, mult_reference):
+    rouge_score, bleu_score, bleu4_score, meteor_score = 0, 0, 0, 0
+    for ref, cand in tqdm(zip(mult_reference, candidates), total=len(mult_reference)):
+        rouge_score += scorer.score(ref, cand)['rougeL'].recall
+        cand_split = wordpunct_tokenize(cand)
+        ref_split = wordpunct_tokenize(ref)
+        bleu4_score += sentence_bleu([ref], cand, weights=(0.0, 0.0, 0.0, 1.0))
+        bleu_score += sentence_bleu([ref], cand)
+        meteor_score += meteor_scorer([ref_split], cand_split)
+    rouge_score, bleu_score, bleu4_score, meteor_score = rouge_score / (len(candidates)), bleu_score / (len(candidates)), bleu4_score / (len(candidates)), meteor_score / (len(candidates))
+    P, R, F1 = score(candidates, mult_reference, lang="en", verbose=True)
+    bert_score = R.mean().item()
+    #print(f"Model: {model_name}")
+    print(f"BLEU Score: {bleu_score}")
+    print(f"BLEU-4 Score: {bleu4_score}")
+    print(f"METEOR Score: {meteor_score}")
+    print(f"ROUGE Score: {rouge_score}")
+    print(f"BERT Score: {bert_score}")
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    parser = transformers.HfArgumentParser((ServeArguments,))
+    serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    dataset_path = "/data/musiccaps/musiccaps_val"
+    ds = load_from_disk(dataset_path)
+    shuffled_ds = ds.shuffle(seed=1234)
+    model, tokenizer = load_trained_lora_model(
+        model_name_or_path=serve_args.model_name_or_path,
+        model_lora_path=serve_args.model_lora_path,
+        load_bits=serve_args.load_bits,
+        use_multi_task=MultiTaskType(serve_args.use_multi_task),
+        tasks_config=serve_args.tasks_config
+    )
+
+    predictions = []
+    references = []
+    content_phrase = random.choice(PRETRAIN_PHRASES)
+#    for data_point_id in range(len(ds)):
+    print("len(ds)", len(ds))
+    for data_point in tqdm(ds):
+        print(data_point["audio"])
+        # data_point = ds[data_point_id]
+        input_json = {"messages": [{"role": "user", "content": content_phrase}], "sounds": [data_point["audio"]]}
+        output_json = generate(input_json)
+        print("Prediction ", output_json["output"])
+        print("Reference ", data_point["caption"])
+        print()
+        print()
+        predictions.append(output_json["output"])
+        references.append(data_point["caption"])
+
+    pairs = {"predictions": predictions, "references": references}
+
+    evaluate(predictions, references)
+
+    with open('/experiments/captioning/mert_tasks_separate_backbone_train_001_ft/checkpoint_1985_test/musiccaps_val_fixed_prompt.yaml', 'w') as file:
+        yaml.dump(pairs, file, default_flow_style=False)

src/sonicverse/scripts/evaluate_mullama.py

@@ -0,0 +1,115 @@
+from dataclasses import dataclass, field
+import logging
+from flask import Flask, request, jsonify
+import transformers
+import torch
+from datasets import load_from_disk
+from multi_token.model_utils import MultiTaskType
+from multi_token.training import ModelArguments
+from multi_token.inference import load_trained_lora_model
+from multi_token.data_tools import encode_chat
+import evaluate
+import random
+import bert_score
+
+PRETRAIN_PHRASES = [
+    "What is happening in the given music <sound>?",
+    "Describe the sound. <sound>",
+    "Describe the music. <sound>",
+    "<sound> Provide a description of the music.",
+    "<sound> Provide a description of the sound.",
+    "Can you interpret <sound>?",
+    "Please explain what's happening in <sound>",
+    "What does <sound> represent?",
+    "Could you describe <sound> for me?",
+    "What's the content of <sound>?",
+    "Can you depict <sound>?",
+    "What is <sound>?",
+    "In the music clip, <sound>, what is happening?",
+    "Provide a description of the music. <sound>",
+    "Provide a description of the sound. <sound>",
+    "Provide a caption for the sound. <sound>",
+    "Provide a caption for the music. <sound>",
+]
+
+@dataclass
+class ServeArguments(ModelArguments):
+    port: int = field(default=8080)
+    host: str = field(default="0.0.0.0")
+    load_bits: int = field(default=16)
+    max_new_tokens: int = field(default=128)
+    temperature: float = field(default=0.01)
+
+def generate(input_json):
+    encoded_dict = encode_chat(input_json, tokenizer, model.modalities)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
+            max_new_tokens=serve_args.max_new_tokens,
+            use_cache=True,
+            do_sample=True,
+            temperature=serve_args.temperature,
+            modality_inputs={
+                m.name: [encoded_dict[m.name]] for m in model.modalities
+            },
+        )
+    outputs = tokenizer.decode(
+        output_ids[0, encoded_dict["input_ids"].shape[0]:],
+        skip_special_tokens=True,
+    ).strip()
+    return {"output": outputs}
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    parser = transformers.HfArgumentParser((ServeArguments,))
+    serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    dataset_path = "/data/musicbench_multitoken_official_split/val"
+    ds = load_from_disk(dataset_path)
+
+    # Load MU-LLaMA model and tokenizer
+    model_name_or_path = "mu-llama/MU-LLaMA"
+    model = transformers.LlamaForCausalLM.from_pretrained(model_name_or_path)
+    tokenizer = transformers.LlamaTokenizer.from_pretrained(model_name_or_path)
+
+    predictions = []
+    references = []
+    content_phrase = random.choice(PRETRAIN_PHRASES)
+    for data_point_id in range(100):
+        data_point = ds[data_point_id]
+        input_json = {"messages": [{"role": "user", "content": content_phrase}], "sounds": data_point["sounds"]}
+        output_json = generate(input_json)
+        print("Prediction ", output_json["output"])
+        print("Reference ", data_point["messages"][1]["content"])
+        print()
+        print()
+        predictions.append(output_json["output"])
+        references.append(data_point["messages"][1]["content"])
+
+    # Load evaluation metrics
+    bleu = evaluate.load("bleu")
+    meteor = evaluate.load("meteor")
+    rouge = evaluate.load("rouge")
+    
+    # Compute BLEU scores
+    bleu_results = bleu.compute(predictions=predictions, references=references, max_order=4)
+#    bleu_score = sum(bleu_results[f"bleu{i}"] for i in range(1, 5)) / 4
+    print(bleu_results)
+
+    # Compute METEOR score
+    meteor_results = meteor.compute(predictions=predictions, references=references)
+    meteor_score = meteor_results["meteor"]
+
+    # Compute ROUGE-L score
+    rouge_results = rouge.compute(predictions=predictions, references=references, rouge_types=["rougeL"])
+    #rouge_l_score = rouge_results["rougeL"].mid.fmeasure
+    print(rouge_results)
+
+    # Compute BERT-Score
+    P, R, F1 = bert_score.score(predictions, references, lang="en", rescale_with_baseline=True)
+    bert_score_f1 = F1.mean().item()
+
+    # Print results
+#    print(f"BLEU Score: {bleu_score}")
+    print(f"METEOR Score: {meteor_score}")
+#    print(f"ROUGE-L Score: {rouge_l_score}")
+    print(f"BERT-Score F1: {bert_score_f1}")

src/sonicverse/scripts/evaluate_temp.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass, field
+import logging
+from flask import Flask, request, jsonify
+import transformers
+import torch
+from datasets import load_from_disk
+from multi_token.model_utils import MultiTaskType
+from multi_token.training import ModelArguments
+from multi_token.inference import load_trained_lora_model
+from multi_token.data_tools import encode_chat
+import evaluate
+import random
+import bert_score
+import os
+
+os.environ['HF_EVALUATE_OFFLINE'] = '1'
+
+PRETRAIN_PHRASES = ["Describe the audio in detail <sound>"]
+
+PRETRAIN_PHRASES_old = [
+    "What is happening in the given music <sound>?",
+    "Describe the sound. <sound>",
+    "Describe the music. <sound>",
+    "<sound> Provide a description of the music.",
+    "<sound> Provide a description of the sound.",
+    "Can you interpret <sound>?",
+    "Please explain what's happening in <sound>",
+    "What does <sound> represent?",
+    "Could you describe <sound> for me?",
+    "What's the content of <sound>?",
+    "Can you depict <sound>?",
+    "What is <sound>?",
+    "In the music clip, <sound>, what is happening?",
+    "Provide a description of the music. <sound>",
+    "Provide a description of the sound. <sound>",
+    "Provide a caption for the sound. <sound>",
+    "Provide a caption for the music. <sound>",
+]
+
+@dataclass
+class ServeArguments(ModelArguments):
+    port: int = field(default=8080)
+    host: str = field(default="0.0.0.0")
+    load_bits: int = field(default=16)
+    max_new_tokens: int = field(default=128)
+    temperature: float = field(default=0.01)
+
+def generate(input_json):
+    encoded_dict = encode_chat(input_json, tokenizer, model.modalities)
+    with torch.inference_mode():
+        output_ids = model.generate(
+            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
+            max_new_tokens=serve_args.max_new_tokens,
+            use_cache=True,
+            do_sample=True,
+            temperature=serve_args.temperature,
+            modality_inputs={
+                m.name: [encoded_dict[m.name]] for m in model.modalities
+            },
+        )
+    outputs = tokenizer.decode(
+        output_ids[0, encoded_dict["input_ids"].shape[0]:],
+        skip_special_tokens=True,
+    ).strip()
+    return {"output": outputs}
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    parser = transformers.HfArgumentParser((ServeArguments,))
+    serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+    dataset_path = "/data/musicbench_multitoken_official_split/val"
+    ds = load_from_disk(dataset_path)
+    model, tokenizer = load_trained_lora_model(
+        model_name_or_path=serve_args.model_name_or_path,
+        model_lora_path=serve_args.model_lora_path,
+        load_bits=serve_args.load_bits,
+        use_multi_task=MultiTaskType(serve_args.use_multi_task),
+        tasks_config=serve_args.tasks_config
+    )
+
+    predictions = []
+    references = []
+    content_phrase = random.choice(PRETRAIN_PHRASES)
+    for data_point_id in range(10):
+        data_point = ds[data_point_id]
+        input_json = {"messages": [{"role": "user", "content": content_phrase}], "sounds": data_point["sounds"]}
+        output_json = generate(input_json)
+        print("Prediction ", output_json["output"])
+        print("Reference ", data_point["messages"][1]["content"])
+        print()
+        print()
+        predictions.append(output_json["output"])
+        references.append(data_point["messages"][1]["content"])
+
+    # Load evaluation metrics
+    bleu = evaluate.load("bleu")
+    meteor = evaluate.load("meteor")
+    rouge = evaluate.load("rouge")
+    
+    # Compute BLEU scores
+    bleu_results = bleu.compute(predictions=predictions, references=references, max_order=4)
+    print(bleu_results)
+    #bleu_score = sum(bleu_results[f"bleu{i}"] for i in range(1, 5)) / 4
+
+    # Compute METEOR score
+    meteor_results = meteor.compute(predictions=predictions, references=references)
+    meteor_score = meteor_results["meteor"]
+
+    # Compute ROUGE-L score
+    rouge_results = rouge.compute(predictions=predictions, references=references, rouge_types=["rougeL"])
+#    rouge_l_score = rouge_results["rougeL"].mid.fmeasure
+    print(rouge_results)
+
+    # Compute BERT-Score
+    P, R, F1 = bert_score.score(predictions, references, lang="en", rescale_with_baseline=True)
+    bert_score_f1 = F1.mean().item()
+
+    # Print results
+    #print(f"BLEU Score: {bleu_score}")
+    print(f"METEOR Score: {meteor_score}")
+#    print(f"ROUGE-L Score: {rouge_l_score}")
+    print(f"BERT-Score F1: {bert_score_f1}")