made classifier.py return probability tensor as well

Author: dyumanaditya

pyreason/scripts/learning/classification/classifier.py

@@ -85,5 +85,5 @@ def forward(self, x, t1: int = 0, t2: int = 0):
             fact_str = f'{self.model_name}({class_name}) : [{lower:.3f}, {upper:.3f}]'
             fact = Fact(fact_str, name=f'{self.model_name}-{class_name}-fact', start_time=t1, end_time=t2)
             facts.append(fact)
-        return output, facts
+        return output, probabilities, facts
 

tests/test_classifier.py

@@ -38,7 +38,7 @@ def test_classifier_integration():
     t2 = 0
 
     # Run the forward pass to get the model output and the corresponding PyReason facts.
-    output, facts = logic_classifier(input_tensor, t1, t2)
+    output, probabilities, facts = logic_classifier(input_tensor, t1, t2)
 
     # Assert that the output is a tensor.
     assert isinstance(output, torch.Tensor), "The model output should be a torch.Tensor"