fix embeddings bug

citoverse · Jan 13, 2025 · bba7594 · bba7594
1 parent 185f86b
commit bba7594
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/R/dnn.R b/R/dnn.R
@@ -242,6 +242,7 @@ dnn <- function(formula = NULL,
   Z_formula = tmp_data$Z_terms
   data = tmp_data$data
 
+
   if(!is.null(Z)) {
 
     Z_args = list()
@@ -250,7 +251,7 @@ dnn <- function(formula = NULL,
     }
 
     embeddings = list(inputs = ncol(Z),
-                      dims = apply(Z, 2, function(z) length(levels(as.factor(z)))), # input embeddings
+                      dims = tmp_data$Zlvls, # input embeddings
                       args = Z_args )
   } else {
     embeddings = NULL
@@ -775,7 +776,10 @@ predict.citodnn <- function(object, newdata = NULL,
     else X <- torch::torch_tensor(stats::model.matrix(stats::as.formula(stats::delete.response(object$call$formula)), data.frame(newdata), xlev = object$data$xlvls)[,-1,drop=FALSE])
 
     if(!is.null(object$model_properties$embeddings)) {
-      tmp = do.call(cbind, lapply(object$Z_formula, function(term) newdata[, term]))
+      tmp = do.call(cbind, lapply(object$Z_formula, function(term) {
+        if(!is.factor(newdata[[term]])) stop(paste0(term, " must be factor."))
+        newdata[[term]] |> as.integer()
+      }) )
       Z = torch::torch_tensor(tmp, dtype = torch::torch_long())
     }
   }

diff --git a/R/utils.R b/R/utils.R
@@ -668,13 +668,19 @@ get_X_Y = function(formula, X, Y, data) {
     out = list(X = X, Y = Y, formula = formula, data = data, Z = NULL, Z_terms = NULL)
   } else {
     terms = sapply(Specials$terms, as.character)
+
+    Zlvls =
+      sapply(terms, function(i) {
+        if(!is.factor(data[,i])) stop("Embeddings must be passed as factor/categorical feature.")
+        return(nlevels(data[,i]))
+      })
     Z =
       lapply(terms, function(i) {
         return(as.integer(data[,i]))
       })
     Z = do.call(cbind, Z)
     colnames(Z) = terms
-    out = list(X = X, Y = Y, formula = formula, data = data, Z = Z, Z_terms = terms, Z_args = Specials$args)
+    out = list(X = X, Y = Y, formula = formula, data = data, Z = Z, Z_terms = terms, Z_args = Specials$args, Zlvls = Zlvls)
   }
   out$old_formula = old_formula
   return(out)