fixes

Author: stephantul

model2vec/train/base.py

@@ -16,7 +16,7 @@
 
 
 class FinetunableStaticModel(nn.Module):
-    def __init__(self, *, vectors: torch.Tensor, tokenizer: Tokenizer, out_dim: int = 2, pad_id: int = 0) -> None:
+    def __init__(self, *, vectors: torch.Tensor, tokenizer: Tokenizer, out_dim: int = 2, pad_id: int = 0, token_mapping: list[int] | None = None) -> None:
         """
         Initialize a trainable StaticModel from a StaticModel.
 
@@ -38,14 +38,19 @@ def __init__(self, *, vectors: torch.Tensor, tokenizer: Tokenizer, out_dim: int
             )
             self.vectors = vectors.float()
 
+        if token_mapping is not None:
+            self.token_mapping = torch.tensor(token_mapping, dtype=torch.int64)
+        else:
+            self.token_mapping = torch.arange(len(vectors), dtype=torch.int64)
+        self.token_mapping = nn.Parameter(self.token_mapping, requires_grad=False)
         self.embeddings = nn.Embedding.from_pretrained(vectors.clone(), freeze=False, padding_idx=pad_id)
         self.head = self.construct_head()
         self.w = self.construct_weights()
         self.tokenizer = tokenizer
 
     def construct_weights(self) -> nn.Parameter:
         """Construct the weights for the model."""
-        weights = torch.zeros(len(self.vectors))
+        weights = torch.zeros(len(self.token_mapping))
         weights[self.pad_id] = -10_000
         return nn.Parameter(weights)
 
@@ -66,11 +71,16 @@ def from_static_model(cls: type[ModelType], *, model: StaticModel, out_dim: int
         """Load the model from a static model."""
         model.embedding = np.nan_to_num(model.embedding)
         embeddings_converted = torch.from_numpy(model.embedding)
+        if model.token_mapping is not None:
+            token_mapping = [i for _, i in sorted(model.token_mapping.items(), key=lambda x: x[0])]
+        else:
+            token_mapping = None
         return cls(
             vectors=embeddings_converted,
             pad_id=model.tokenizer.token_to_id("[PAD]"),
             out_dim=out_dim,
             tokenizer=model.tokenizer,
+            token_mapping=token_mapping,
             **kwargs,
         )
 
@@ -90,7 +100,8 @@ def _encode(self, input_ids: torch.Tensor) -> torch.Tensor:
         w = w * zeros
         # Add a small epsilon to avoid division by zero
         length = zeros.sum(1) + 1e-16
-        embedded = self.embeddings(input_ids)
+        input_ids_embeddings = self.token_mapping[input_ids]
+        embedded = self.embeddings(input_ids_embeddings)
         # Weigh each token
         embedded = torch.bmm(w[:, None, :], embedded).squeeze(1)
         # Mean pooling by dividing by the length
@@ -118,16 +129,17 @@ def tokenize(self, texts: list[str], max_length: int | None = 512) -> torch.Tens
         return pad_sequence(encoded_ids, batch_first=True, padding_value=self.pad_id)
 
     @property
-    def device(self) -> str:
+    def device(self) -> torch.device:
         """Get the device of the model."""
         return self.embeddings.weight.device
 
     def to_static_model(self) -> StaticModel:
         """Convert the model to a static model."""
         emb = self.embeddings.weight.detach().cpu().numpy()
         w = torch.sigmoid(self.w).detach().cpu().numpy()
+        token_mapping = {i: int(token_id) for i, token_id in enumerate(self.token_mapping.tolist())}
 
-        return StaticModel(emb * w[:, None], self.tokenizer, normalize=True)
+        return StaticModel(vectors=emb, weights=w, tokenizer=self.tokenizer, normalize=True, token_mapping=token_mapping)
 
 
 class TextDataset(Dataset):

model2vec/train/classifier.py

@@ -38,6 +38,7 @@ def __init__(
         hidden_dim: int = 512,
         out_dim: int = 2,
         pad_id: int = 0,
+        token_mapping: list[int] | None = None,
     ) -> None:
         """Initialize a standard classifier model."""
         self.n_layers = n_layers
@@ -46,7 +47,7 @@ def __init__(
         self.classes_: list[str] = [str(x) for x in range(out_dim)]
         # multilabel flag will be set based on the type of `y` passed to fit.
         self.multilabel: bool = False
-        super().__init__(vectors=vectors, out_dim=out_dim, pad_id=pad_id, tokenizer=tokenizer)
+        super().__init__(vectors=vectors, out_dim=out_dim, pad_id=pad_id, tokenizer=tokenizer, token_mapping=token_mapping)
 
     @property
     def classes(self) -> np.ndarray: