multiclass fix

Author: SeBorgey

autointent/modules/scoring/_gcn/gcn_scorer.py

@@ -92,7 +92,8 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
         self._label_embedder = Embedder(self.label_embedder_config)
 
         x_tensor = torch.tensor(self._embedder.embed(utterances, TaskTypeEnum.classification))
-        y_tensor = torch.tensor(labels, dtype=torch.float)
+        y_tensor_dtype = torch.float if self._multilabel else torch.long
+        y_tensor = torch.tensor(labels, dtype=y_tensor_dtype)
 
         intent_texts = [f"intent {i}" for i in range(self._n_classes)]
         self._label_embeddings = torch.tensor(
@@ -107,14 +108,16 @@ def fit(self, utterances: list[str], labels: ListOfLabels) -> None:
             p_reweight=self.p_reweight,
             tau_threshold=self.tau_threshold,
         )
-        self._model.set_correlation_matrix(y_tensor)
-        self._train_model(x_tensor, y_tensor)
 
-    def _train_model(self, train_x: torch.Tensor, train_y: torch.Tensor) -> None:
+        y_corr_tensor = y_tensor if self._multilabel else torch.nn.functional.one_hot(y_tensor, self._n_classes)
+        self._model.set_correlation_matrix(y_corr_tensor.float())
+
+        criterion = nn.BCEWithLogitsLoss() if self._multilabel else nn.CrossEntropyLoss()
+        self._train_model(x_tensor, y_tensor, criterion)
+
+    def _train_model(self, train_x: torch.Tensor, train_y: torch.Tensor, criterion: nn.Module) -> None:
         train_dataset = TensorDataset(train_x, train_y)
         train_dataloader = DataLoader(train_dataset, batch_size=self.torch_config.batch_size, shuffle=True)
-
-        criterion = nn.BCEWithLogitsLoss()
         optimizer = torch.optim.Adam(self._model.parameters(), lr=self.torch_config.learning_rate)
 
         self._model.to(self.torch_config.device)
@@ -142,7 +145,10 @@ def predict(self, utterances: list[str]) -> npt.NDArray[Any]:
             for i in range(0, len(x_tensor), self.torch_config.batch_size):
                 batch_x = x_tensor[i : i + self.torch_config.batch_size].to(self.torch_config.device)
                 outputs = self._model(batch_x, self._label_embeddings)
-                probs = torch.sigmoid(outputs).cpu().numpy()
+                if self._multilabel:
+                    probs = torch.sigmoid(outputs).cpu().numpy()
+                else:
+                    probs = torch.softmax(outputs, dim=1).cpu().numpy()
                 all_probs.append(probs)
 
         return np.concatenate(all_probs, axis=0)

tests/modules/scoring/test_gcn_scorer.py

@@ -23,36 +23,57 @@ def multilabel_dataset():
     return Dataset.from_dict(data)
 
 
-def test_gcn_scorer_fit_predict(multilabel_dataset):
-    scorer = GCNScorer(
-        embedder_config="prajjwal1/bert-tiny",
-        num_train_epochs=1,
-        batch_size=2,
-    )
+@pytest.fixture
+def multiclass_dataset():
+    data = {
+        "train": [
+            {"utterance": "utterance 1", "label": 0},
+            {"utterance": "utterance 2", "label": 1},
+            {"utterance": "utterance 3", "label": 2},
+            {"utterance": "utterance 4", "label": 0},
+        ],
+        "intents": [
+            {"id": 0, "name": "intent_0"},
+            {"id": 1, "name": "intent_1"},
+            {"id": 2, "name": "intent_2"},
+        ],
+    }
+    return Dataset.from_dict(data)
+
+
+def test_gcn_scorer_multilabel(multilabel_dataset):
+    scorer = GCNScorer(embedder_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=2)
     train_utterances = multilabel_dataset["train"]["utterance"]
     train_labels = multilabel_dataset["train"]["label"]
-
     scorer.fit(train_utterances, train_labels)
+    test_utterances = ["test 1", "test 2"]
+    predictions = scorer.predict(test_utterances)
 
-    test_utterances = ["test utterance 1", "test utterance 2"]
+    assert isinstance(predictions, np.ndarray)
+    assert predictions.shape == (2, 3)
+    assert np.all((predictions >= 0) & (predictions <= 1))
+
+
+def test_gcn_scorer_multiclass(multiclass_dataset):
+    scorer = GCNScorer(embedder_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=2)
+    train_utterances = multiclass_dataset["train"]["utterance"]
+    train_labels = multiclass_dataset["train"]["label"]
+    scorer.fit(train_utterances, train_labels)
+    test_utterances = ["test 1", "test 2"]
     predictions = scorer.predict(test_utterances)
 
     assert isinstance(predictions, np.ndarray)
     assert predictions.shape == (2, 3)
     assert np.all((predictions >= 0) & (predictions <= 1))
+    np.testing.assert_allclose(predictions.sum(axis=1), 1.0, atol=1e-6)
 
 
 def test_gcn_scorer_dump_load(tmp_path, multilabel_dataset):
-    scorer = GCNScorer(
-        embedder_config="prajjwal1/bert-tiny",
-        num_train_epochs=1,
-        batch_size=2,
-    )
+    scorer = GCNScorer(embedder_config="prajjwal1/bert-tiny", num_train_epochs=1, batch_size=2)
     train_utterances = multilabel_dataset["train"]["utterance"]
     train_labels = multilabel_dataset["train"]["label"]
     scorer.fit(train_utterances, train_labels)
-
-    test_utterances = ["test utterance 1", "test utterance 2"]
+    test_utterances = ["test utterance 1"]
     original_predictions = scorer.predict(test_utterances)
 
     scorer.dump(str(tmp_path))