diff --git a/model2vec/train/base.py b/model2vec/train/base.py index a60a150..9c5eaa1 100644 --- a/model2vec/train/base.py +++ b/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): diff --git a/tests/test_trainable.py b/tests/test_trainable.py index f8b2af2..fb61d73 100644 --- a/tests/test_trainable.py +++ b/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: