add trait, target_in_fit

Author: ablaom

src/MLJModelInterface.jl

@@ -8,6 +8,7 @@ const MODEL_TRAITS = [
     :predict_scitype,
     :transform_scitype,
     :inverse_transform_scitype,
+    :target_in_fit,
     :is_pure_julia,
     :package_name,
     :package_license,

src/model_traits.jl

@@ -32,6 +32,9 @@ end
 StatTraits.is_supervised(::Type{<:Supervised}) = true
 StatTraits.is_supervised(::Type{<:SupervisedAnnotator}) = true
 
+StatTraits.target_in_fit(::Type{<:Supervised}) = true
+StatTraits.target_in_fit(::Type{<:Unsupervised}) = false
+
 StatTraits.prediction_type(::Type{<:Deterministic}) = :deterministic
 StatTraits.prediction_type(::Type{<:Probabilistic}) = :probabilistic
 StatTraits.prediction_type(::Type{<:Interval}) = :interval
@@ -73,7 +76,15 @@ function supervised_fit_data_scitype(M)
     return ret
 end
 
-StatTraits.fit_data_scitype(M::Type{<:Unsupervised}) = Tuple{input_scitype(M)}
+# helper to determine the scitype of unsupervised models
+function unsupervised_fit_data_scitype(M)
+    I = input_scitype(M)
+    T = target_scitype(M)
+    target_in_fit(M) && return Tuple{I, T}
+    return Tuple{I}
+end
+
+StatTraits.fit_data_scitype(M::Type{<:Unsupervised}) = unsupervised_fit_data_scitype(M)
 StatTraits.fit_data_scitype(::Type{<:Static}) = Tuple{}
 StatTraits.fit_data_scitype(M::Type{<:Supervised}) = supervised_fit_data_scitype(M)
 

test/model_traits.jl

@@ -23,6 +23,10 @@ end
 @mlj_model mutable struct UA <: UnsupervisedAnnotator
 end
 
+@mlj_model mutable struct SupervisedTransformer <: Unsupervised
+end
+
+
 foo(::P1) = 0
 bar(::P1) = nothing
 
@@ -34,6 +38,10 @@ M.package_name(::Type{<:U1}) = "Bach"
 M.package_url(::Type{<:U1}) = "www.did_he_write_565.com"
 M.human_name(::Type{<:U1}) = "funky model"
 
+M.target_in_fit(::Type{<:SupervisedTransformer}) = true
+M.target_scitype(::Type{<:SupervisedTransformer}) = Continuous
+M.input_scitype(::Type{<:SupervisedTransformer}) = Finite
+
 @testset "traits" begin
     ms = S1()
     mu = U1(a=42, b=sin)
@@ -42,6 +50,7 @@ M.human_name(::Type{<:U1}) = "funky model"
     mi = I1()
     sa = SA()
     ua = UA()
+    supervised_transformer = SupervisedTransformer()
 
     @test input_scitype(ms)  == Unknown
     @test output_scitype(ms) == Unknown
@@ -115,6 +124,11 @@ M.human_name(::Type{<:U1}) = "funky model"
     setfull()
 
     @test Set(implemented_methods(mp)) == Set([:clean!,:bar,:foo])
+
+    @test fit_data_scitype(mu) == Tuple{Unknown};;;
+    @test fit_data_scitype(mu) == Tuple{Unknown}
+    @test fit_data_scitype(supervised_transformer) == Tuple{Finite,Continuous}
+ 
 end
 
 @testset "`_density` - helper for predict_scitype fallback" begin