Upload folder using huggingface_hub

cosa/compute_embeddings.py

@@ -0,0 +1,148 @@
+import os
+import argparse
+import torch
+from tqdm import tqdm
+from transformers import (
+    AutoTokenizer, AutoModel,
+    BertTokenizer, BertModel,
+    CLIPTokenizer, CLIPTextModel,
+    T5Tokenizer, T5EncoderModel
+)
+
+import sys
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "osm_clip")))
+from model import OSMBind
+
+
+def average_pool(last_hidden_states, attention_mask):
+    """Computes average pooling of hidden states, masking padding tokens."""
+    masked_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return masked_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+
+def get_tokenizer_and_model(encoder_type='bert', checkpoint_path=None, taglist_path = None, tagvocab_path = None):
+    if encoder_type == 'bert':
+        model_name = 'bert-base-uncased'
+        tokenizer = BertTokenizer.from_pretrained(model_name)
+        model = BertModel.from_pretrained(model_name)
+        embedding_fn = lambda outputs, batch_dict: outputs.pooler_output.squeeze()
+
+    elif encoder_type == 'clip':
+        model_name = 'openai/clip-vit-large-patch14'
+        tokenizer = CLIPTokenizer.from_pretrained(model_name)
+        model = CLIPTextModel.from_pretrained(model_name)
+
+        def clip_embedding_fn(outputs, batch_dict):
+            input_ids = batch_dict['input_ids']
+            eos_token_id = tokenizer.eos_token_id
+            seq_lengths = (input_ids == eos_token_id).nonzero(as_tuple=True)[1]
+
+            embeddings = []
+            for i in range(input_ids.size(0)):
+                eos_pos = seq_lengths[i] if i < len(seq_lengths) else (input_ids[i] != tokenizer.pad_token_id).sum() - 1
+                embeddings.append(outputs.last_hidden_state[i, eos_pos, :])
+            return torch.stack(embeddings)
+
+        embedding_fn = clip_embedding_fn
+
+    elif encoder_type == 'e5':
+        model_name = 'intfloat/e5-base-v2'
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModel.from_pretrained(model_name)
+        embedding_fn = lambda outputs, batch_dict: average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+
+    elif encoder_type == 't5':
+        model_name = 't5-base'
+        tokenizer = T5Tokenizer.from_pretrained(model_name)
+        model = T5EncoderModel.from_pretrained(model_name)
+        embedding_fn = lambda outputs, batch_dict: average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
+
+    elif 'osm' in encoder_type:
+        text_backbone = encoder_type.split('-')[1] if '-' in encoder_type else 'clip'
+        model = OSMBind(taglist_path=taglist_path, tagvocab_path=tagvocab_path, text_backbone=text_backbone)
+        ckpt = torch.load(checkpoint_path, map_location='cpu')
+        model.load_state_dict(ckpt['state_dict'], strict=False)
+        model.eval().cuda()
+        tokenizer = None
+
+        def osm_embedding_fn(outputs, batch_dict):
+            return model.text_encoder.encode_batch(batch_dict['sentences'])
+
+        embedding_fn = osm_embedding_fn
+
+    else:
+        raise ValueError(f"Unsupported encoder_type: {encoder_type}")
+
+    model.eval()
+    return tokenizer, model, embedding_fn
+
+
+def generate_embeddings(taglist_path, tag_vocab_path, output_path,
+                                     encoder_type='bert', checkpoint_path=None):
+    # Load taglist and vocab
+    taglist = torch.load(taglist_path, weights_only = True)  # list of tuples of tag indices
+    tag_vocab = torch.load(tag_vocab_path, weights_only = True)
+    tag_index = {v: k for k, v in tag_vocab.items()}  # index -> tag string
+
+    # Convert taglist tuples to "sentences" of tag strings
+    sentences = []
+    for tl in taglist:
+        words = [tag_index[idx] for idx in tl]
+        sentences.append(" ".join(words))
+
+    # Optional prompt formatting
+    if encoder_type == 'e5':
+        sentences = [f"query: {s}" for s in sentences]
+    elif encoder_type == 't5':
+        sentences = [f"embedding: {s}" for s in sentences]
+
+    # Load model
+    tokenizer, model, embedding_fn = get_tokenizer_and_model(encoder_type, checkpoint_path, taglist_path = taglist_path, tagvocab_path = tag_vocab_path)
+    device = next(model.parameters()).device if hasattr(model, 'parameters') else torch.device('cpu')
+
+    # Generate embeddings
+    embeddings = []
+    print("Encoding taglists...")
+    for sentence in tqdm(sentences):
+        if 'osm' in encoder_type:
+            batch_dict = {'sentences': [sentence]}
+            outputs = None
+        else:
+            inputs = tokenizer([sentence], return_tensors='pt', padding=True, truncation=True)
+            batch_dict = {k: v.to(device) for k, v in inputs.items()}
+            outputs = model(**batch_dict)
+
+        with torch.inference_mode():
+            emb = embedding_fn(outputs, batch_dict)
+            if emb.ndim == 1:
+                emb = emb.unsqueeze(0)
+            embeddings.append(emb.cpu())
+
+    embeddings = torch.cat(embeddings, dim=0)
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    torch.save(embeddings, output_path)
+    print(f"Saved {len(sentences)} taglist embeddings to {output_path}")
+
+
+# ========================
+# Command Line Interface
+# ========================
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate embeddings for taglists")
+    parser.add_argument("--taglist_path", type=str, required=True, help="Path to taglist_vocab.pt")
+    parser.add_argument("--tag_vocab_path", type=str, required=True, help="Path to tag_vocab.pt")
+    parser.add_argument("--output_path", type=str, required=True, help="Path to save embeddings tensor")
+    parser.add_argument("--encoder_type", type=str,
+                        choices=["bert", "clip", "e5", "t5", "osm-clip", "osm-e5", "osm-bert"],
+                        default="bert")
+    parser.add_argument("--checkpoint_path", type=str, default=None, help="Optional checkpoint for OSMBind")
+
+    args = parser.parse_args()
+
+    generate_embeddings(
+        taglist_path=args.taglist_path,
+        tag_vocab_path=args.tag_vocab_path,
+        output_path=args.output_path,
+        encoder_type=args.encoder_type,
+        checkpoint_path=args.checkpoint_path
+    )

cosa/cosa.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:981a8ec6c089d019dbe54afd34693d3617db8b28837cf5adf013702563b6f73a
+size 2365975368

cosa/model.py

@@ -0,0 +1,290 @@
+import torch
+import torch.nn as nn
+import os
+import numpy as np
+import torch.nn.functional as F
+import pytorch_lightning as pl
+from datasets import OSMDataset
+from torch.utils.data import DataLoader
+import random
+from typing import Optional, List, Tuple, Literal
+from image_encoder import SatlasPretrainEncoder
+from text_encoder import TextEncoder
+from orthogonal_adamw import OrthogonalAdamW
+from configs.config_e5 import config
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from lightning.pytorch.loggers import WandbLogger
+from utils import generate_tag_poly_pairs
+import matplotlib.pyplot as plt
+import io
+import wandb
+from PIL import Image
+
+
+# This performs a typical InfoNCE loss
+def contrastive_loss(image_feats: torch.Tensor, text_feats: torch.Tensor, logit_scale: torch.Tensor) -> torch.Tensor:
+    logits = torch.matmul(image_feats, text_feats.t()) * logit_scale
+    labels = torch.arange(logits.size(0), device=logits.device)
+
+    return F.cross_entropy(logits, labels), logits
+
+
+class OSMBind(pl.LightningModule):
+    def __init__(self, train_dataset=None, val_dataset=None, **kwargs):
+        super().__init__()
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+
+        self.image_encoder = SatlasPretrainEncoder(fpn=True, model_name="Aerial_SwinB_SI",
+                                                   out_dim=768, num_extra_fpn_layers=4)
+        taglist_vocab = torch.load(kwargs.get("taglist_path"), weights_only = True)
+        tag_vocab_inverted = torch.load(kwargs.get("tagvocab_path"), weights_only = True) # str -> int
+        tag_vocab = {v: k for k, v in tag_vocab_inverted.items()} # int -> str
+        self.text_encoder = TextEncoder(taglist_vocab, tag_vocab, 
+                                        model_name=kwargs.get("text_backbone"))
+        # for param in self.text_encoder.parameters():
+        #     param.requires_grad = False
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) # softer scale for misaligned encoders
+    
+        self.batch_size = kwargs.get("batch_size")
+        self.num_workers = kwargs.get("num_workers")
+        self.lr = kwargs.get("lr", 1e-4)
+        self.num_samples = kwargs.get("num_samples") # number of OSM classes sampled
+        self.ort_grad = kwargs.get("ort_grad")
+    
+    def forward(self, sat_img: torch.Tensor, pixel_tensor: torch.Tensor):
+        full_image_feats = self.image_encoder(sat_img)  # [B, D, H', W']
+        sampled_tag_tensor, image_poly_feats = generate_tag_poly_pairs(pixel_tensor, full_image_feats, K=self.num_samples) # [K], [K, D]
+        text_sampled_feats = self.text_encoder(sampled_tag_tensor)  # [K, D]
+        
+        return image_poly_feats, text_sampled_feats # [K, D], [K, D]
+    
+    def shared_step(self, batch):
+        sat_img, pixel_tensor = batch
+        image_poly_feats, text_sampled_feats = self(sat_img, pixel_tensor)  # [K, D], [K, D]
+
+        # contrastive loss for whole batch
+        image_feats_norm = F.normalize(image_poly_feats, dim=1)
+        text_feats_norm = F.normalize(text_sampled_feats, dim=1)
+        logit_scale = self.logit_scale.exp()
+        loss, logits = contrastive_loss(image_feats_norm, text_feats_norm,
+                                            logit_scale=logit_scale)
+        return loss, logits
+    
+    def log_similarity_matrix(self, logits):
+        mat = logits.detach().cpu().numpy()
+        fig, ax = plt.subplots(figsize=(6,6))
+        cax = ax.matshow(mat, cmap="viridis")
+        fig.colorbar(cax)
+        ax.set_xlabel("Text samples")
+        ax.set_ylabel("Image samples")
+        ax.set_title("Similarity Matrix")
+
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png')
+        buf.seek(0)
+        plt.close(fig)
+
+        # ✅ Fix: Convert buffer to PIL Image
+        image = Image.open(buf)
+
+        if isinstance(self.logger, WandbLogger):
+            self.logger.experiment.log({
+                "similarity_matrix": wandb.Image(image),
+                "global_step": self.global_step
+            })
+
+    def training_step(self, batch, batch_idx):
+        loss, logits = self.shared_step(batch)
+        self.log('train_loss', loss, sync_dist=True, prog_bar=True, on_epoch=True, batch_size=self.batch_size)
+        self.log('temperature', self.logit_scale.exp().item(), prog_bar=True, on_epoch=True)
+        if self.global_step % 500 == 0:
+            self.log_similarity_matrix(logits)
+        # Log histogram of similarity scores every step
+        if self.logger and hasattr(self.logger.experiment, "log"):
+            self.logger.experiment.log({"logits_hist": wandb.Histogram(logits.detach().cpu().numpy())})
+
+        # Optionally log mean and max of logits for monitoring
+        self.log("logits_mean", logits.mean(), on_step=True, on_epoch=False, prog_bar=True)
+        self.log("logits_max", logits.max(), on_step=True, on_epoch=False, prog_bar=True)
+        return loss
+    
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        min_log_scale = np.log(1 / 1.0)    
+        max_log_scale = np.log(1 / 0.01)   
+        self.logit_scale.data.clamp_(min_log_scale, max_log_scale)
+
+    def on_after_backward(self):
+        if self.global_rank == 0 and self.current_epoch == 0:
+            for name, param in self.named_parameters():
+                if param.requires_grad and param.grad is None:
+                    print(f"⚠️ Unused parameter: {name}")
+
+    def validation_step(self, batch, batch_idx):
+        loss, _ = self.shared_step(batch)
+        self.log('val_loss', loss, sync_dist=True, prog_bar=True, on_epoch=True, batch_size=self.batch_size)
+        return loss
+
+    def train_dataloader(self):
+        if self.train_dataset is None:
+            raise ValueError("This model was initialized without a training dataset.")
+        return DataLoader(self.train_dataset,
+                          batch_size=self.batch_size,
+                          num_workers=self.num_workers,
+                          shuffle=True,
+                          persistent_workers=False)
+
+    def val_dataloader(self):
+        if self.val_dataset is None:
+            raise ValueError("This model was initialized without a validation dataset.")
+        return DataLoader(self.val_dataset,
+                          batch_size=self.batch_size,
+                          num_workers=self.num_workers,
+                          shuffle=False,
+                          persistent_workers=False)
+
+    def configure_optimizers(self):
+        params = self.parameters()
+        if self.ort_grad:
+            self.optim = OrthogonalAdamW(
+                params,
+                lr=self.lr,
+                betas=(0.9, 0.98),
+                beta_ort=0.9,
+                eps=1e-6,
+                weight_decay=0.01
+            )
+        else:
+            self.optim = torch.optim.AdamW(
+                params,
+                lr=self.lr,
+                betas=(0.9, 0.98),
+                eps=1e-6,
+                weight_decay=0.01
+            )
+        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer=self.optim,
+            T_0=20
+        )
+
+        return [self.optim], [self.scheduler]
+    
+    def sim_map_inf(self, sat_image: torch.Tensor, raw_text: str) -> torch.Tensor:
+        """
+        Args:
+            sat_image: [1, 3, 512, 512] tensor (already normalized)
+            raw_text: str, e.g., "building"
+
+        Returns:
+            sim_map: [512, 512] similarity map between image and text embedding
+        """
+        assert sat_image.dim() == 4 and sat_image.size(0) == 1, "Expected input of shape [1, 3, H, W]"
+
+        # Step 1: Extract spatial features
+        with torch.no_grad():
+            # image features
+            feat_map = self.image_encoder(sat_image)  # [1, D, H', W']
+            feat_map = feat_map.squeeze(0)            # [D, H', W']
+            feat_map_upsampled = F.interpolate(feat_map.unsqueeze(0), size=(512, 512), mode='bilinear', align_corners=False).squeeze(0)  # [D, 512, 512]
+            feat_map_upsampled = F.normalize(feat_map_upsampled, dim=0)  # [D, 512, 512]
+
+            # text features
+            text_feat = self.text_encoder.encode_raw_text(raw_text)
+
+            # cosine sim
+            text_feat = F.normalize(text_feat, dim=0)
+            feat_map_upsampled = F.normalize(feat_map_upsampled, dim=0)
+            sim_map = torch.einsum('chw,c->hw', feat_map_upsampled, text_feat)  # [512, 512]
+
+        return sim_map
+
+    def encode_text(self, text: str) -> torch.Tensor:
+        with torch.no_grad():
+            return self.text_encoder.encode_raw_text(text)
+    
+    def encode_image(self, image: torch.Tensor) -> torch.Tensor:
+        with torch.no_grad():
+            return self.image_encoder(image)
+
+def seed_everything(seed=42):
+    """
+    seed: int
+    """
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+if __name__=='__main__':
+    import warnings
+    warnings.filterwarnings("ignore")
+    torch.set_warn_always(False)
+
+    seed_everything()
+    train_dataset = OSMDataset(metadata_path = config.train_csv,
+                               image_dir=config.sat_img_dir, 
+                               pixel_tensor_dir=config.pixel_tensors_dir,
+                               mode='train')
+    val_dataset = OSMDataset(metadata_path = config.val_csv,
+                               image_dir=config.sat_img_dir, 
+                               pixel_tensor_dir=config.pixel_tensors_dir,
+                               mode='val')
+    
+    # from torch.utils.data import Subset
+    # train_dataset = Subset(train_dataset, range(1000))
+    # val_dataset = Subset(val_dataset, range(200))
+    
+    kwargs = {
+        'batch_size':config.batch_size, 
+        'num_workers': config.num_workers,
+        'num_samples': config.num_contrastive_samples,
+        'ort_grad': config.ort_grad,
+        'lr': config.lr,
+        'taglist_vocab_path': config.taglist_vocab_path,
+        'tag_vocab_path': config.tag_vocab_path,
+        'text_backbone': config.text_backbone
+    }
+    
+    model = OSMBind(train_dataset, val_dataset, **kwargs)
+    torch.cuda.empty_cache()
+
+    checkpoint_path = '/data/b.j.wei/rendersynth/osm_clip/checkpoints/osmclip_e5/osmclip_config_e5-epoch=39-val_loss=3.23.ckpt'
+    if checkpoint_path:
+        ckpt = torch.load(checkpoint_path, map_location='cpu')
+        model.load_state_dict(ckpt['state_dict'])
+
+    checkpoint = ModelCheckpoint(
+        monitor='val_loss',
+        dirpath=config.save_dir,
+        filename=config.filename,
+        mode='min',
+        save_top_k=1,
+        every_n_epochs=1
+    )
+
+    early_stop_callback = EarlyStopping(
+        monitor='val_loss',
+        patience=15,
+        mode='min'
+    )
+
+    logger = WandbLogger(project="osmclip", 
+                         name=f"{config.experiment_name}")
+
+    trainer = pl.Trainer(
+        accelerator='gpu',
+        devices=config.devices,
+        strategy='ddp',
+        max_epochs=config.max_epochs,
+        num_nodes=1,
+        callbacks=[checkpoint, early_stop_callback],
+        accumulate_grad_batches=config.accumulate_grad_batches,
+        log_every_n_steps=5, 
+        logger = logger #wandb logger
+        )
+    
+    trainer.fit(model)

cosa/text_encoder.py

@@ -0,0 +1,374 @@
+import torch
+from transformers import (
+    AutoTokenizer, AutoModel, 
+    BertTokenizer, BertModel, 
+    CLIPTokenizer, CLIPTextModel
+)
+import torch.nn as nn
+import pytorch_lightning as pl
+from typing import List
+from abc import ABC, abstractmethod
+import random
+
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+def taglist_index_to_sentence(taglist_vocab, tag_vocab, taglist_indices, subsample: bool = True):
+    """
+    Convert a tensor or list of taglist indices to a list of tag sentences.
+    Optionally, randomly shuffle and sample a subset of tags for each sentence.
+    
+    Args:
+        taglist_vocab: List of tuples of tag IDs.
+        tag_vocab: Dictionary mapping tag ID to tag string.
+        taglist_indices: Tensor or list of indices into taglist_vocab.
+        seed: Random seed for reproducibility.
+        subsample: If True, randomly subsample tags in each sentence.
+    
+    Returns:
+        tag_sentences: List of strings (tag sentences).
+    """
+    if isinstance(taglist_indices, torch.Tensor):
+        taglist_indices = taglist_indices.view(-1).tolist()
+
+    tag_sentences = []
+    
+    for idx in taglist_indices:
+        tag_ids = taglist_vocab[idx]
+        tags = [tag_vocab[tid].lower().replace('=', ' ') for tid in tag_ids]
+        
+        if subsample and len(tags) > 1:
+            n_sample = random.randint(1, len(tags))  # Choose how many tags to keep
+            tags = random.sample(tags, n_sample)     # Sample without replacement
+
+        random.shuffle(tags)  # Randomize order
+        sentence = ' '.join(tags)
+        tag_sentences.append(sentence)
+
+    return tag_sentences
+
+
+def average_pool(last_hidden_states, attention_mask):
+    masked_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return masked_hidden.sum(dim=1) / attention_mask.sum(dim=1, keepdim=True)
+
+
+class BaseTextEncoder(nn.Module, ABC):
+    def __init__(self, model_name: str):
+        super().__init__()
+        self.model_name = model_name
+        self.tokenizer = None
+        self.model = None
+        self.embedding_dim = None
+
+    @abstractmethod
+    def encode(self, sentences: List[str], device: str = 'cpu') -> torch.Tensor:
+        """
+        Encode a list of sentences into a tensor of embeddings.
+        Must be implemented by subclasses.
+        """
+        pass
+
+class BertTextEncoder(BaseTextEncoder):
+    def __init__(self, model_name='bert-base-uncased'):
+        super().__init__(model_name)
+        self.tokenizer = BertTokenizer.from_pretrained(model_name)
+        self.model = BertModel.from_pretrained(model_name)
+        self.embedding_dim = self.model.config.hidden_size
+
+    def encode(self, sentences, device='cpu'):
+        self.model.to(device)
+        inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+        return self.model(**inputs).pooler_output
+
+
+class CLIPTextEncoder(BaseTextEncoder):
+    def __init__(self, model_name='openai/clip-vit-large-patch14', local_tokenizer_path=None):
+        super().__init__(model_name)
+        local_tokenizer_path = "/u/cherd/.cache/huggingface/hub/models--openai--clip-vit-large-patch14/snapshots/32bd64288804d66eefd0ccbe215aa642df71cc41"
+
+        if local_tokenizer_path is not None:
+            self.tokenizer = CLIPTokenizer.from_pretrained(local_tokenizer_path)
+            self.model = CLIPTextModel.from_pretrained(local_tokenizer_path)
+        else:
+            self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
+            self.model = CLIPTextModel.from_pretrained(model_name, from_flax=True)
+        self.embedding_dim = self.model.config.hidden_size
+
+    def encode(self, sentences, device='cpu'):
+        self.model.to(device)
+        inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+        input_ids = inputs['input_ids']
+        eos_token_id = self.tokenizer.eos_token_id
+        pad_token_id = self.tokenizer.pad_token_id
+
+        outputs = self.model(**inputs)
+        last_hidden = outputs.last_hidden_state  # [B, T, D]
+
+        batch_size = input_ids.size(0)
+        embeddings = []
+
+        for i in range(batch_size):
+            input_seq = input_ids[i]
+            eos_positions = (input_seq == eos_token_id).nonzero(as_tuple=True)[0]
+
+            if len(eos_positions) > 0:
+                eos_idx = eos_positions[-1]  # take last EOS (safe for duplicates)
+            else:
+                eos_idx = (input_seq != pad_token_id).sum() - 1  # fallback to last non-padding token
+
+            embeddings.append(last_hidden[i, eos_idx, :])
+
+        return torch.stack(embeddings)
+
+class E5TextEncoder(BaseTextEncoder):
+    def __init__(self, model_name='intfloat/e5-base'):
+        super().__init__(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+        self.model.pooler = None
+        self.embedding_dim = self.model.config.hidden_size
+
+    def encode(self, sentences, device='cpu'):
+        self.model.to(device)
+        sentences = [f"query: {s}" for s in sentences]  # official prompt for e5 (for features as per documentation)
+        inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+        outputs = self.model(**inputs)
+        return average_pool(outputs.last_hidden_state, inputs['attention_mask'])
+
+class GritLMTextEncoder(BaseTextEncoder):
+    def __init__(self, model_name='nomic-ai/nomic-bert-base-punc'):
+        super().__init__(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+        self.embedding_dim = self.model.config.hidden_size
+        self.proj_head = nn.Linear(self.embedding_dim, 768) # to match other encoders
+
+    def encode(self, sentences, device='cpu'):
+        self.model.to(device)
+        inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+        outputs = self.model(**inputs)
+        pooled = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
+        return self.proj_head(pooled)
+
+
+class TextEncoder(pl.LightningModule):
+    def __init__(self, taglist_vocab: List[tuple], tag_vocab: dict, model_name='bert'):
+        super().__init__()
+        self.taglist_vocab = taglist_vocab
+        self.tag_vocab = tag_vocab
+
+        model_name = model_name.lower()
+        encoder_map = {
+            'bert': lambda: BertTextEncoder('bert-base-uncased'),
+            'clip': lambda: CLIPTextEncoder('openai/clip-vit-large-patch14'),
+            'e5': lambda: E5TextEncoder('intfloat/e5-base'),
+            'gritlm': lambda: GritLMTextEncoder('nomic-ai/nomic-bert-base-punc')
+        }
+
+        if model_name not in encoder_map:
+            raise ValueError(f"Unsupported model_name: {model_name}. Choose from {list(encoder_map.keys())}")
+        print(f"Text backbone: {model_name}")
+        self.encoder = encoder_map[model_name]()  # Instantiate the selected encoder
+        # self.embedding_dim = 768
+
+    def forward(self, taglist_tensor: torch.Tensor) -> torch.Tensor:
+        tag_indices = taglist_tensor.tolist()
+        tag_sentences = taglist_index_to_sentence(self.taglist_vocab, self.tag_vocab, tag_indices, subsample=True) # randomize subsampling tags
+        embeddings = self.encoder.encode(tag_sentences, device=self.device)
+        return embeddings
+
+    def encode_raw_text(self, raw_text: str) -> torch.Tensor:
+        """
+        Encode a single raw string into an embedding for queries
+        """
+        return self.encoder.encode([raw_text], device=self.device)[0] 
+    
+    def encode_batch(self, raw_texts: List[str]) -> torch.Tensor:
+        """
+        Encode a batch of raw strings into embeddings for queries
+        """
+        return self.encoder.encode(raw_texts, device=self.device)
+
+
+# import torch
+# from transformers import (
+#     AutoTokenizer, AutoModel, 
+#     BertTokenizer, BertModel, 
+#     CLIPTokenizer, CLIPTextModel
+# )
+# import torch.nn as nn
+# import pytorch_lightning as pl
+# from typing import List
+# from abc import ABC, abstractmethod
+# import random
+
+# import os
+# os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+# def taglist_index_to_sentence(taglist_vocab, tag_vocab, taglist_indices, subsample: bool = True):
+#     """
+#     Convert a tensor or list of taglist indices to a list of tag sentences.
+#     Optionally, randomly shuffle and sample a subset of tags for each sentence.
+    
+#     Args:
+#         taglist_vocab: List of tuples of tag IDs.
+#         tag_vocab: Dictionary mapping tag ID to tag string.
+#         taglist_indices: Tensor or list of indices into taglist_vocab.
+#         seed: Random seed for reproducibility.
+#         subsample: If True, randomly subsample tags in each sentence.
+    
+#     Returns:
+#         tag_sentences: List of strings (tag sentences).
+#     """
+#     if isinstance(taglist_indices, torch.Tensor):
+#         taglist_indices = taglist_indices.view(-1).tolist()
+
+#     tag_sentences = []
+    
+#     for idx in taglist_indices:
+#         tag_ids = taglist_vocab[idx]
+#         tags = [tag_vocab[tid].lower().replace('=', ' ') for tid in tag_ids]
+        
+#         if subsample and len(tags) > 1:
+#             n_sample = random.randint(1, len(tags))  # Choose how many tags to keep
+#             tags = random.sample(tags, n_sample)     # Sample without replacement
+
+#         random.shuffle(tags)  # Randomize order
+#         sentence = ' '.join(tags)
+#         tag_sentences.append(sentence)
+
+#     return tag_sentences
+
+
+# def average_pool(last_hidden_states, attention_mask):
+#     masked_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+#     return masked_hidden.sum(dim=1) / attention_mask.sum(dim=1, keepdim=True)
+
+
+# class BaseTextEncoder(nn.Module, ABC):
+#     def __init__(self, model_name: str):
+#         super().__init__()
+#         self.model_name = model_name
+#         self.tokenizer = None
+#         self.model = None
+#         self.embedding_dim = None
+
+#     @abstractmethod
+#     def encode(self, sentences: List[str], device: str = 'cpu') -> torch.Tensor:
+#         """
+#         Encode a list of sentences into a tensor of embeddings.
+#         Must be implemented by subclasses.
+#         """
+#         pass
+
+# class BertTextEncoder(BaseTextEncoder):
+#     def __init__(self, model_name='bert-base-uncased'):
+#         super().__init__(model_name)
+#         self.tokenizer = BertTokenizer.from_pretrained(model_name)
+#         self.model = BertModel.from_pretrained(model_name)
+#         self.embedding_dim = self.model.config.hidden_size
+
+#     def encode(self, sentences, device='cpu'):
+#         self.model.to(device)
+#         inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+#         return self.model(**inputs).pooler_output
+
+
+# class CLIPTextEncoder(BaseTextEncoder):
+#     def __init__(self, model_name='openai/clip-vit-large-patch14'):
+#         super().__init__(model_name)
+#         self.tokenizer = CLIPTokenizer.from_pretrained(model_name)
+#         self.model = CLIPTextModel.from_pretrained(model_name)
+#         self.embedding_dim = self.model.config.hidden_size
+
+#     def encode(self, sentences, device='cpu'):
+#         self.model.to(device)
+#         inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+#         input_ids = inputs['input_ids']
+#         eos_token_id = self.tokenizer.eos_token_id
+#         pad_token_id = self.tokenizer.pad_token_id
+
+#         outputs = self.model(**inputs)
+#         last_hidden = outputs.last_hidden_state  # [B, T, D]
+
+#         batch_size = input_ids.size(0)
+#         embeddings = []
+
+#         for i in range(batch_size):
+#             input_seq = input_ids[i]
+#             eos_positions = (input_seq == eos_token_id).nonzero(as_tuple=True)[0]
+
+#             if len(eos_positions) > 0:
+#                 eos_idx = eos_positions[-1]  # take last EOS (safe for duplicates)
+#             else:
+#                 eos_idx = (input_seq != pad_token_id).sum() - 1  # fallback to last non-padding token
+
+#             embeddings.append(last_hidden[i, eos_idx, :])
+
+#         return torch.stack(embeddings)
+
+# class E5TextEncoder(BaseTextEncoder):
+#     def __init__(self, model_name='intfloat/e5-base'):
+#         super().__init__(model_name)
+#         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+#         self.model = AutoModel.from_pretrained(model_name)
+#         self.model.pooler = None
+#         self.embedding_dim = self.model.config.hidden_size
+
+#     def encode(self, sentences, device='cpu'):
+#         self.model.to(device)
+#         sentences = [f"query: {s}" for s in sentences]  # official prompt for e5 (for features as per documentation)
+#         inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+#         outputs = self.model(**inputs)
+#         return average_pool(outputs.last_hidden_state, inputs['attention_mask'])
+
+# class GritLMTextEncoder(BaseTextEncoder):
+#     def __init__(self, model_name='nomic-ai/nomic-bert-base-punc'):
+#         super().__init__(model_name)
+#         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+#         self.model = AutoModel.from_pretrained(model_name)
+#         self.embedding_dim = self.model.config.hidden_size
+#         self.proj_head = nn.Linear(self.embedding_dim, 768) # to match other encoders
+
+#     def encode(self, sentences, device='cpu'):
+#         self.model.to(device)
+#         inputs = self.tokenizer(sentences, return_tensors='pt', padding=True, truncation=True).to(device)
+#         outputs = self.model(**inputs)
+#         pooled = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
+#         return self.proj_head(pooled)
+
+
+# class TextEncoder(pl.LightningModule):
+#     def __init__(self, taglist_vocab: List[tuple], tag_vocab: dict, model_name='bert'):
+#         super().__init__()
+#         self.taglist_vocab = taglist_vocab
+#         self.tag_vocab = tag_vocab
+
+#         model_name = model_name.lower()
+#         encoder_map = {
+#             'bert': lambda: BertTextEncoder('bert-base-uncased'),
+#             'clip': lambda: CLIPTextEncoder('openai/clip-vit-large-patch14'),
+#             'e5': lambda: E5TextEncoder('intfloat/e5-base'),
+#             'gritlm': lambda: GritLMTextEncoder('nomic-ai/nomic-bert-base-punc')
+#         }
+
+#         if model_name not in encoder_map:
+#             raise ValueError(f"Unsupported model_name: {model_name}. Choose from {list(encoder_map.keys())}")
+#         print(f"Text backbone: {model_name}")
+#         self.encoder = encoder_map[model_name]()  # Instantiate the selected encoder
+#         # self.embedding_dim = 768
+
+#     def forward(self, taglist_tensor: torch.Tensor) -> torch.Tensor:
+#         tag_indices = taglist_tensor.tolist()
+#         tag_sentences = taglist_index_to_sentence(self.taglist_vocab, self.tag_vocab, tag_indices, subsample=True) # randomize subsampling tags
+#         embeddings = self.encoder.encode(tag_sentences, device=self.device)
+#         return embeddings
+
+#     def encode_raw_text(self, raw_text: str) -> torch.Tensor:
+#         """
+#         Encode a single raw string into an embedding for queries
+#         """
+#         return self.encoder.encode([raw_text], device=self.device)[0]