Add docstrings

Author: eliascarv

ext/StatsLearnModelsMLJModelInterfaceExt.jl

@@ -10,20 +10,23 @@ isprobabilistic(model::MI.Probabilistic) = true
 function SLM.fit(model::MI.Model, input, output)
   cols = Tables.columns(output)
   names = Tables.columnnames(cols)
-  y = Tables.getcolumn(cols, first(names))
+  target = first(names)
+  y = Tables.getcolumn(cols, target)
   data = MI.reformat(model, input, y)
   fitresult, _... = MI.fit(model, 0, data...)
-  SLM.FittedModel(model, fitresult)
+  SLM.FittedModel(model, (fitresult, target))
 end
 
 function SLM.predict(fmodel::SLM.FittedModel{<:MI.Model}, table)
-  (; model, fitresult) = fmodel
+  (; model, cache) = fmodel
+  fitresult, target = cache
   data = MI.reformat(model, table)
-  if isprobabilistic(model)
+  ŷ = if isprobabilistic(model)
     MI.predict_mode(model, fitresult, data...)
   else
     MI.predict(model, fitresult, data...)
   end
+  (; target => ŷ)
 end
 
 end

src/StatsLearnModels.jl

@@ -3,25 +3,29 @@ module StatsLearnModels
 """
     StatsLearnModels.fit(model, input, output) -> FittedModel
 
-TODO
+Fit statistical learning `model` using features in `input` table
+and targets in `output` table. Returns a fitted model with all
+the necessary information for prediction with the `predict` function.
 """
 function fit end
 
 """
     StatsLearnModels.predict(model::FittedModel, table)
 
-TODO
+Predict the target values using the fitted statistical learning `model`
+and a new `table` of features.  
 """
 function predict end
 
 """
-    StatsLearnModels.FittedModel(model, fitresult)
+    StatsLearnModels.FittedModel(model, cache)
 
-TODO
+Wrapper type used to save learning model and auxiliary
+variables needed for prediction.
 """
-struct FittedModel{M,F}
+struct FittedModel{M,C}
   model::M
-  fitresult::F
+  cache::C
 end
 
 end