Error while creating workflow from a recipie using linear models in R

I am training a linear regression model predicting salary from company size (company_size_number) and country (country) using the StackOverflow data.

What I perform is:

Read the data. Split the data into a training set (75%) and a test set (25%).
Create a recipe that converts company_size_number into a factor variable and then transforms the two predictors into dummy variables.
Create the model specification.
Create a workflow object and add the recipe and model specification to it, then fit the model on the training set.
Calculate R² on the test set.

This is my code

library(tidyverse)
library(tidymodels)

so <- read_rds("stackoverflow.rds") 

set.seed(123)
init_split <- initial_split(so)
so_training <- training(init_split)
so_testing <- testing(init_split)

rec <- recipe(salary ~ ., data = so_training %>% select(salary, company_size_number, country)) %>%
  step_num2factor(company_size_number = factor(company_size_number)) %>%
  step_dummy(country, company_size_number)

model_spec <- linear_reg() %>%
  set_engine("lm") %>%
  set_mode("regression")

fit <- workflow() %>%
  add_model(model_spec) %>%
  add_recipe(rec) %>%
  fit(data = so_training)

predict(fit, new_data = so_testing) %>%
  mutate(truth = so_testing$salary) %>%
  rmse(estimate = .pred, truth = truth)

But not able to proceed due to an error:

Error: Please provide a character vector of appropriate length for `levels`.

I presume I am messing up something here in the spec_*()

rec <- recipe(salary ~ ., data = so_training %>% select(salary, company_size_number, country)) %>%
  step_novel(company_size_number = factor(company_size_number)) %>%
  step_dummy(country, company_size_number)

But not sure if this correct. Any inputs would be helpful.

> dput(head(so))
structure(list(country = structure(c(5L, 5L, 4L, 4L, 5L, 5L), .Label = c("Canada", 
"Germany", "India", "United Kingdom", "United States"), class = "factor"), 
    salary = c(63750, 93000, 40625, 45000, 1e+05, 170000), years_coded_job = c(4L, 
    9L, 8L, 3L, 8L, 12L), open_source = c(0, 1, 1, 1, 0, 1), 
    hobby = c(1, 1, 1, 0, 1, 1), company_size_number = c(20, 
    1000, 10000, 1, 10, 100), remote = structure(c(1L, 1L, 1L, 
    1L, 1L, 1L), .Label = c("Remote", "Not remote"), class = "factor"), 
    career_satisfaction = c(8L, 8L, 5L, 10L, 8L, 10L), data_scientist = c(0, 
    0, 1, 0, 0, 0), database_administrator = c(1, 0, 1, 0, 0, 
    0), desktop_applications_developer = c(1, 0, 1, 0, 0, 0), 
    developer_with_stats_math_background = c(0, 0, 0, 0, 0, 0
    ), dev_ops = c(0, 0, 0, 0, 0, 1), embedded_developer = c(0, 
    0, 0, 0, 0, 0), graphic_designer = c(0, 0, 0, 0, 0, 0), graphics_programming = c(0, 
    0, 0, 0, 0, 0), machine_learning_specialist = c(0, 0, 0, 
    0, 0, 0), mobile_developer = c(0, 1, 0, 0, 1, 0), quality_assurance_engineer = c(0, 
    0, 0, 0, 0, 0), systems_administrator = c(1, 0, 1, 0, 0, 
    1), web_developer = c(0, 0, 0, 1, 1, 1)), row.names = c(NA, 
-6L), class = c("tbl_df", "tbl", "data.frame"))

Solution

I have a couple of recommendations for adjustments in what you are doing.

The first is to do the selecting of variables before splitting, so that when you use a formula like salary ~ ., you and/or the functions don't get confused about what is there.
The second is to not use step_num2factor() in the way you have; it would take a lot to get it to work correctly and I think you're better served converting it to a factor before you split. Take a look at this step's documentation to see a more appropriate use for this recipe step, and notice that you have to give it levels. This is the reason you saw the error you did, but honestly I wouldn't try to find the right levels and input them there; I'd do it before splitting.

library(tidyverse)
library(tidymodels)

data("stackoverflow", package = "modeldata")
so <- janitor::clean_names(stackoverflow)

set.seed(123)
init_split <- so %>%
   select(salary, company_size_number, country) %>%
   mutate(company_size_number = factor(company_size_number)) %>%
   initial_split()
so_training <- training(init_split)
so_testing <- testing(init_split)

rec <- recipe(salary ~ ., data = so_training) %>%
   step_dummy(country, company_size_number)

model_spec <- linear_reg() %>%
   set_engine("lm") %>%
   set_mode("regression")

fit <- workflow() %>%
   add_model(model_spec) %>%
   add_recipe(rec) %>%
   fit(data = so_training)

predict(fit, new_data = so_testing) %>%
   mutate(truth = so_testing$salary) %>%
   rmse(estimate = .pred, truth = truth)
#> # A tibble: 1 x 3
#>   .metric .estimator .estimate
#>   <chr>   <chr>          <dbl>
#> 1 rmse    standard      27822.

^{Created on 2021-05-25 by the reprex package (v2.0.0)}