Frozen Embeddings with InversionFromLogitsModel is incorrect

[themurtazanazir]

in tokenize_data.py

def embed_dataset_batch(model: InversionModel, batch: Dict) -> Dict:
    assert "input_ids" in batch.keys(), f"invalid keys {batch.keys()}"
    assert hasattr(model, "call_embedding_model")

    input_ids = batch["input_ids"]
    inputs_str = model.tokenizer.batch_decode(input_ids, skip_special_tokens=True)
    emb_input_ids = model.embedder_tokenizer(
        inputs_str,
        max_length=model.config.max_seq_length,
        truncation=True,
        padding="max_length",
        return_tensors="pt",
    ).to(next(model.parameters()).device)

    with torch.no_grad():
        batch["frozen_embeddings"] = model.call_embedding_model(**emb_input_ids)
    return batch

the tokens of embedder are sent to call_embedding_model.

But in models/inversion_from_logits.py

    def call_embedding_model(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        embedder = self.embedder

        inputs_str = self.tokenizer.batch_decode(input_ids, skip_special_tokens=True)
        emb_input_ids = self.embedder_tokenizer(
            inputs_str,
            max_length=self.config.max_seq_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        ).to(next(self.parameters()).device)

        model_output = embedder(**emb_input_ids)
        return self._process_embedder_output(model_output, emb_input_ids.attention_mask)

This function expects the model.tokenizer's, not model.embedder_tokenizers's tokens.

This causes gibberish tokens to be sent to embedder.

[themurtazanazir]

I can raise a PR if needed.

[jxmorris12]

@themurtazanazir thank you for finding this – a pull request would be amazing!