feat: remove loggit mapping from training

model2vec/train/base.py

@@ -17,7 +17,6 @@
 from model2vec.train.trainer import MetricsFn, default_metrics, resolve_device, run_training_loop
 from model2vec.train.utils import (
     get_probable_pad_token_id,
-    logit,
     to_pipeline,
     train_test_split,
 )
@@ -101,10 +100,10 @@ def _remove_unk(self, token_ids: list[int]) -> list[int]:
     def construct_weights(self) -> nn.Parameter:
         """Construct the weights for the model."""
         if self._weights is not None:
-            w = logit(self._weights)
+            w = self._weights
         else:
-            w = torch.zeros(len(self.token_mapping)).float()
-            w[self.pad_id] = -10_000
+            w = torch.ones(len(self.token_mapping)).float()
+            w[self.pad_id] = 0
         return nn.Parameter(w, requires_grad=not self.freeze_weights)
 
     def construct_head(self) -> nn.Sequential:
@@ -216,7 +215,6 @@ def _encode(self, input_ids: torch.Tensor) -> torch.Tensor:
         embedded = self.embeddings(input_ids_embeddings)
 
         w = self.w[input_ids]
-        w = torch.sigmoid(w)
         w = w * zeros
         # Weigh each token
         embedded = torch.bmm(w[:, None, :], embedded).squeeze(1)
@@ -270,7 +268,7 @@ def to_static_model(self) -> StaticModel:
         with torch.no_grad():
             emb = self.embeddings.weight
             emb = emb.cpu().numpy()
-            w = torch.sigmoid(self.w).cpu().numpy()
+            w = self.w.cpu().numpy()
 
         # If the weights and emb are the same length, the model was not quantized before training.
         if len(w) == len(emb):

tests/test_trainable.py

@@ -88,7 +88,7 @@ def test_init_classifier_from_model_w(mock_vectors: np.ndarray, mock_tokenizer:
     assert torch.all(s._weights == torch.ones(len(mock_vectors)))
     w = s.construct_weights()
     assert w.shape[0] == mock_vectors.shape[0]
-    assert torch.all(w == logit(torch.ones(len(mock_vectors))))
+    assert torch.all(w == torch.ones(len(mock_vectors)))
 
 
 def test_pad_token(mock_tokenizer: Tokenizer) -> None: