"""Code to determine embeddings for text and images."""
import warnings
from typing import List, Union
import numpy as np
import torch
from nilearn import datasets
from nilearn.image import concat_imgs, load_img, new_img_like, resample_to_img
from nilearn.maskers import NiftiMapsMasker, SurfaceMapsMasker
from nimare.dataset import Dataset
from nimare.meta.kernel import MKDAKernel
from tqdm import tqdm
from braindec.utils import _get_device, _vol_surfimg, images_have_same_fov
def _coordinates_to_image(dset: Dataset, kernel: str = "mkda"):
if kernel == "mkda":
kernel = MKDAKernel()
else:
raise ValueError(f"Kernel {kernel} not supported.")
return kernel.transform(dset, return_type="image")
[docs]
class TextEmbedding:
def __init__(
self,
model_name: str = "BrainGPT/BrainGPT-7B-v0.2",
max_length: int = None,
batch_size: int = 1,
device: str = None,
):
"""
Initialize the embedding generator with specified model and parameters.
Args:
model_name: Name of the model to use. Supported models are:
- "mistralai/Mistral-7B-v0.1"
- "meta-llama/Llama-2-7b-chat-hf"
- "BrainGPT/BrainGPT-7B-v0.1"
- "BrainGPT/BrainGPT-7B-v0.2"
max_length: Maximum token length for each chunk
batch_size: Batch size for processing (total number of papers to process at once)
device: Device to use for computation. If None, the device is automatically selected.
"""
self.device = _get_device() if device is None else device
self.model_name = model_name
self.batch_size = batch_size
if model_name == "mistralai/Mistral-7B-v0.1":
from transformers import AutoModel, AutoTokenizer
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModel.from_pretrained(model_name).to(self.device)
self.max_length = 8192 if max_length is None else max_length
elif model_name == "meta-llama/Llama-2-7b-chat-hf":
from transformers import AutoModelForCausalLM, AutoTokenizer
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device)
self.max_length = 4096 if max_length is None else max_length
elif model_name == "BrainGPT/BrainGPT-7B-v0.1":
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
config = PeftConfig.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
self.model = PeftModel.from_pretrained(model, model_name).to(self.device)
self.tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
self.max_length = 4096 if max_length is None else max_length
elif model_name == "BrainGPT/BrainGPT-7B-v0.2":
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
config = PeftConfig.from_pretrained(model_name)
# The config file has path to the base model instead of the model name
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
self.model = PeftModel.from_pretrained(model, model_name).to(self.device)
self.tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
self.max_length = 8192 if max_length is None else max_length
else:
raise ValueError(f"Model name {model_name} not supported.")
self.model.eval()
self.tokenizer.pad_token = self.tokenizer.eos_token
[docs]
def clear_device_cache(self):
"""Clear memory cache for the current device type."""
if self.device == "cuda":
torch.cuda.empty_cache()
elif self.device == "mps":
# MPS (Apple Silicon) garbage collection
torch.mps.empty_cache()
[docs]
def mean_pooling(self, token_embeddings, attention_mask):
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
sentence_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
input_mask_expanded.sum(1), min=1e-9
)
return sentence_embeddings
[docs]
def generate_embedding(self, token_embeddings, attention_mask) -> np.ndarray:
"""
Generate embedding from token embeddings.
Args:
token_embeddings: Token embeddings
attention_mask: Attention mask
Returns:
Numpy array containing the average embedding
"""
embeddings = self.mean_pooling(token_embeddings, attention_mask)
return embeddings.cpu().numpy()
[docs]
def get_token_embeddings(self, tokenized: dict) -> np.ndarray:
"""
Get token embeddings for a single text chunk.
Args:
tokenized: Tokenized dictionary containing input_ids and attention_mask
Returns:
Numpy array containing the embedding
"""
# Send to device
tokenized = {k: v.to(self.device) for k, v in tokenized.items()}
# Generate embeddings
with torch.no_grad():
outputs = self.model(**tokenized, output_hidden_states=True)
return outputs.hidden_states[-1]
[docs]
def chunk_text(self, texts: str) -> List[str]:
"""
Split text into chunks that respect the model's token limit.
Args:
text: Input text to be chunked
Returns:
List of text chunks (tokenized dictionaries)
"""
tokenized = self.tokenizer(texts, return_tensors="pt", padding=True)
input_ids = tokenized["input_ids"]
attention_mask = tokenized["attention_mask"]
# Split into chunks of dictionaries
chunks = []
for i in range(0, input_ids.size(1), self.max_length):
chunks.append(
{
"input_ids": input_ids[:, i : i + self.max_length],
"attention_mask": attention_mask[:, i : i + self.max_length],
}
)
return chunks
[docs]
def process_text(self, text: List[str]) -> np.ndarray:
"""
Process text by chunking and averaging token embeddings.
Args:
text: Batch of texts
Returns:
Averaged embedding vector for the entire batch of texts
"""
# Split text into chunks
chunks = self.chunk_text(text)
# Generate embeddings for each chunk
token_embeddings = []
attention_masks = []
for chunk in tqdm(chunks, desc="Processing chunks", leave=False):
token_embeddings.append(self.get_token_embeddings(chunk))
attention_masks.append(chunk["attention_mask"].to(self.device))
# Delete chunks to free up memory
del chunks
token_embeddings = torch.cat(token_embeddings, dim=1)
attention_masks = torch.cat(attention_masks, dim=1)
return self.generate_embedding(token_embeddings, attention_masks)
def __call__(self, texts: Union[str, List[str]]) -> np.ndarray:
"""
Generate embeddings for input text(s).
Args:
text: Input text or list of texts
Returns:
Numpy array of embeddings
"""
if isinstance(texts, str):
return self.process_text([texts])
else:
# Process multiple texts in batches
embeddings = []
for t in tqdm(range(0, len(texts), self.batch_size), desc="Processing batches"):
embeddings.append(self.process_text(texts[t : t + self.batch_size]))
return np.concatenate(embeddings)
[docs]
class ImageEmbedding:
_default_density = {"fsLR": "32k", "fsaverage": "164k", "civet": "32k"}
def __init__(
self,
standardize: bool = False,
nilearn_dir: str = None,
neuromaps_dir: str = None,
atlas: str = "difumo",
dimension: int = 512,
space: str = "MNI152",
density: str = None,
):
"""
Initialize the image embedding generator with specified model.
Args:
model_name: Name of the DeiT model to use
"""
self.nilearn_dir = nilearn_dir
self.neuromaps_dir = neuromaps_dir
self.atlas = atlas
self.dimension = dimension
self.space = space
self.density = density
self._maps_cache = {}
if self.atlas == "difumo":
difumo_kwargs = {
"dimension": self.dimension,
"resolution_mm": 2,
"data_dir": self.nilearn_dir,
}
try:
difumo = datasets.fetch_atlas_difumo(
legacy_format=False,
**difumo_kwargs,
)
except TypeError:
difumo = datasets.fetch_atlas_difumo(**difumo_kwargs)
atlas_maps = difumo.maps
self.atlas_maps = load_img(atlas_maps)
else:
# Implement other atlases
raise ValueError(f"Atlas {atlas} not supported.")
if self.space == "MNI152":
self.masker = NiftiMapsMasker(maps_img=atlas_maps, standardize=standardize)
elif self.space in ["fsLR", "fsaverage", "civet"]:
warnings.warn("Do not use this for now. As the training was done in MNI space.")
self.density = self._default_density[space] if self.density is None else self.density
# Trasnform atlas to surface
atlas_surf = _vol_surfimg(
atlas_maps,
space=self.space,
density=self.density,
neuromaps_dir=self.neuromaps_dir,
)
self.masker = SurfaceMapsMasker(maps_img=atlas_surf, standardize=standardize)
else:
raise ValueError(f"Space {self.space} not supported.")
[docs]
def generate_embedding(self, images) -> np.ndarray:
"""
Generate embedding for a single image.
Args:
image: Input image as a numpy array
Returns:
Numpy array containing the embedding
"""
if self.space == "MNI152":
return self._generate_volume_embedding(images)
if isinstance(images, list):
images = [self._sanitize_image(image) for image in images]
# Concat images to improve performance
images = concat_imgs(images)
else:
images = self._sanitize_image(images)
embeddings = self.masker.fit_transform(images)
if embeddings.ndim == 1:
embeddings = embeddings[None, :]
return embeddings
@staticmethod
def _sanitize_image(image):
image = load_img(image)
image_data = image.get_fdata()
if np.isfinite(image_data).all():
return image
warnings.warn("Non-finite values detected in image data. Replacing them with zeros.")
image_data = np.nan_to_num(image_data, nan=0.0, posinf=0.0, neginf=0.0)
return new_img_like(image, image_data, copy_header=True)
def _get_maps_data(self, reference_img):
if reference_img is None:
cache_key = ("native",)
else:
cache_key = (
tuple(reference_img.shape[:3]),
tuple(reference_img.affine.ravel()),
)
if cache_key not in self._maps_cache:
atlas_img = self.atlas_maps
if reference_img is not None and not images_have_same_fov(atlas_img, reference_img):
atlas_img = resample_to_img(atlas_img, reference_img, interpolation="continuous")
maps_data = atlas_img.get_fdata(dtype=np.float32)
maps_data = np.nan_to_num(maps_data, nan=0.0, posinf=0.0, neginf=0.0)
maps_gram = np.tensordot(
maps_data,
maps_data,
axes=([0, 1, 2], [0, 1, 2]),
).astype(np.float32)
maps_gram += np.eye(maps_gram.shape[0], dtype=np.float32) * 1e-6
self._maps_cache[cache_key] = (maps_data, maps_gram)
return self._maps_cache[cache_key]
def _generate_volume_embedding(self, images) -> np.ndarray:
image = self._sanitize_image(images)
if not images_have_same_fov(image, self.atlas_maps):
image = resample_to_img(image, self.atlas_maps, interpolation="continuous")
image_data = image.get_fdata(dtype=np.float32)
image_data = np.nan_to_num(image_data, nan=0.0, posinf=0.0, neginf=0.0)
if image_data.ndim == 3:
image_data = image_data[..., None]
maps_data, maps_gram = self._get_maps_data(None)
xty = np.tensordot(maps_data, image_data, axes=([0, 1, 2], [0, 1, 2])).astype(np.float32)
embeddings = np.linalg.solve(maps_gram, xty).T
if embeddings.ndim == 1:
embeddings = embeddings[None, :]
return embeddings
def __call__(self, images) -> np.ndarray:
"""
Generate embeddings for input images.
Args:
images: List of input images as numpy arrays
Returns:
Numpy array of embeddings
"""
# Accept nifti and path to image as input
return self.generate_embedding(images)