Adding RoBertaSentenceEmbeddings and XlmRoBertaSentenceEmbeddings annotators

Author: dkincaid

NAMESPACE

@@ -285,6 +285,7 @@ export(nlp_relation_extraction_dl)
 export(nlp_relation_extraction_dl_pretrained)
 export(nlp_relation_extraction_pretrained)
 export(nlp_roberta_embeddings_pretrained)
+export(nlp_roberta_sentence_embeddings_pretrained)
 export(nlp_sentence_detector)
 export(nlp_sentence_detector_dl)
 export(nlp_sentence_detector_dl_pretrained)
@@ -318,6 +319,7 @@ export(nlp_word_embeddings)
 export(nlp_word_embeddings_model)
 export(nlp_word_embeddings_pretrained)
 export(nlp_xlm_roberta_embeddings_pretrained)
+export(nlp_xlm_roberta_sentence_embeddings_pretrained)
 export(nlp_xlnet_embeddings_pretrained)
 export(nlp_yake_model)
 export(set_nlp_version)

R/roberta_sentence_embeddings.R

@@ -0,0 +1,61 @@
+#' Load a pretrained Spark NLP RoBertaSentenceEmbeddings model
+#' 
+#' Create a pretrained Spark NLP \code{RoBertaSentenceEmbeddings} model.
+#' Sentence-level embeddings using RoBERTa. The RoBERTa model was proposed in 
+#' RoBERTa: A Robustly Optimized BERT Pretraining Approach by Yinhan Liu, Myle Ott, 
+#' Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, 
+#' Luke Zettlemoyer, Veselin Stoyanov. It is based on Google's BERT model 
+#' released in 2018.
+#' 
+#' It builds on BERT and modifies key hyperparameters, removing the next-sentence pretraining objective and training with much larger mini-batches and learning rates.
+#' See \url{https://nlp.johnsnowlabs.com/docs/en/annotators#robertabertsentenceembeddings}
+#' 
+#' @template roxlate-pretrained-params
+#' @template roxlate-inputs-output-params
+#' @param batch_size batch size
+#' @param case_sensitive whether to lowercase tokens or not
+#' @param dimension defines the output layer of BERT when calculating embeddings
+#' @param max_sentence_length max sentence length to process
+#' 
+#' @export
+nlp_roberta_sentence_embeddings_pretrained <- function(sc, input_cols, output_col, case_sensitive = NULL,
+                                              batch_size = NULL, dimension = NULL, 
+                                              max_sentence_length = NULL, 
+                                              name = NULL, lang = NULL, remote_loc = NULL) {
+  args <- list(
+    input_cols = input_cols,
+    output_col = output_col,
+    case_sensitive = case_sensitive,
+    batch_size = batch_size,
+    dimension = dimension,
+    max_sentence_length = max_sentence_length
+  ) %>%
+    validator_nlp_roberta_sentence_embeddings()
+  
+  model_class <- "com.johnsnowlabs.nlp.embeddings.RoBertaSentenceEmbeddings"
+  model <- pretrained_model(sc, model_class, name, lang, remote_loc)
+  spark_jobj(model) %>%
+    sparklyr::jobj_set_param("setInputCols", args[["input_cols"]]) %>% 
+    sparklyr::jobj_set_param("setOutputCol", args[["output_col"]]) %>% 
+    sparklyr::jobj_set_param("setCaseSensitive", args[["case_sensitive"]]) %>% 
+    sparklyr::jobj_set_param("setBatchSize", args[["batch_size"]]) %>% 
+    sparklyr::jobj_set_param("setDimension", args[["dimension"]]) %>% 
+    sparklyr::jobj_set_param("setMaxSentenceLength", args[["max_sentence_length"]])
+  
+  new_nlp_roberta_sentence_embeddings(model)
+}
+
+#' @import forge
+validator_nlp_roberta_sentence_embeddings <- function(args) {
+  args[["input_cols"]] <- cast_string_list(args[["input_cols"]])
+  args[["output_col"]] <- cast_string(args[["output_col"]])
+  args[["batch_size"]] <- cast_nullable_integer(args[["batch_size"]])
+  args[["case_sensitive"]] <- cast_nullable_logical(args[["case_sensitive"]])
+  args[["dimension"]] <- cast_nullable_integer(args[["dimension"]])
+  args[["max_sentence_length"]] <- cast_nullable_integer(args[["max_sentence_length"]])
+  args
+}
+
+new_nlp_roberta_sentence_embeddings <- function(jobj) {
+  sparklyr::new_ml_transformer(jobj, class = "nlp_roberta_sentence_embeddings")
+}

R/xlm_roberta_sentence_embeddings.R

@@ -0,0 +1,61 @@
+#' Load a pretrained Spark NLP XlmRoBertaSentenceEmbeddings model
+#' 
+#' Create a pretrained Spark NLP \code{XlmRoBertaSentenceEmbeddings} model.
+#' See \url{https://nlp.johnsnowlabs.com/docs/en/annotators#xlmrobertasentenceembeddings}
+#' 
+#' Sentence-level embeddings using XLM-RoBERTa. The XLM-RoBERTa model was proposed in 
+#' Unsupervised Cross-lingual Representation Learning at Scale by Alexis Conneau, 
+#' Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán, 
+#' Edouard Grave, Myle Ott, Luke Zettlemoyer and Veselin Stoyanov. It is based on 
+#' Facebook's RoBERTa model released in 2019. It is a large multi-lingual language model, 
+#' trained on 2.5TB of filtered CommonCrawl data.
+#' 
+#' @template roxlate-pretrained-params
+#' @template roxlate-inputs-output-params
+#' @param batch_size batch size
+#' @param case_sensitive whether to lowercase tokens or not
+#' @param dimension defines the output layer of BERT when calculating embeddings
+#' @param max_sentence_length max sentence length to process
+#' 
+#' @export
+nlp_xlm_roberta_sentence_embeddings_pretrained <- function(sc, input_cols, output_col, case_sensitive = NULL,
+                                              batch_size = NULL, dimension = NULL, 
+                                              max_sentence_length = NULL, 
+                                              name = NULL, lang = NULL, remote_loc = NULL) {
+  args <- list(
+    input_cols = input_cols,
+    output_col = output_col,
+    case_sensitive = case_sensitive,
+    batch_size = batch_size,
+    dimension = dimension,
+    max_sentence_length = max_sentence_length
+  ) %>%
+    validator_nlp_xlm_roberta_sentence_embeddings()
+  
+  model_class <- "com.johnsnowlabs.nlp.embeddings.XlmRoBertaSentenceEmbeddings"
+  model <- pretrained_model(sc, model_class, name, lang, remote_loc)
+  spark_jobj(model) %>%
+    sparklyr::jobj_set_param("setInputCols", args[["input_cols"]]) %>% 
+    sparklyr::jobj_set_param("setOutputCol", args[["output_col"]]) %>% 
+    sparklyr::jobj_set_param("setCaseSensitive", args[["case_sensitive"]]) %>% 
+    sparklyr::jobj_set_param("setBatchSize", args[["batch_size"]]) %>% 
+    sparklyr::jobj_set_param("setDimension", args[["dimension"]]) %>% 
+    sparklyr::jobj_set_param("setMaxSentenceLength", args[["max_sentence_length"]])
+  
+  new_nlp_xlm_roberta_sentence_embeddings(model)
+}
+
+#' @import forge
+validator_nlp_xlm_roberta_sentence_embeddings <- function(args) {
+  args[["input_cols"]] <- cast_string_list(args[["input_cols"]])
+  args[["output_col"]] <- cast_string(args[["output_col"]])
+  args[["batch_size"]] <- cast_nullable_integer(args[["batch_size"]])
+  args[["case_sensitive"]] <- cast_nullable_logical(args[["case_sensitive"]])
+  args[["dimension"]] <- cast_nullable_integer(args[["dimension"]])
+  args[["max_sentence_length"]] <- cast_nullable_integer(args[["max_sentence_length"]])
+  args
+}
+
+new_nlp_xlm_roberta_sentence_embeddings <- function(jobj) {
+  sparklyr::new_ml_transformer(jobj, class = "nlp_xlm_roberta_sentence_embeddings")
+}

inst/sparkml/class_mapping.json

@@ -79,6 +79,8 @@
   "com.johnsnowlabs.nlp.embeddings.AlbertEmbeddings": "nlp_albert_embeddings",
   "com.johnsnowlabs.nlp.embeddings.BertEmbeddings": "nlp_bert_embeddings",
   "com.johnsnowlabs.nlp.embeddings.BertSentenceEmbeddings": "nlp_bert_sentence_embeddings",
+  "com.johnsnowlabs.nlp.embeddings.RoBertaSentenceEmbeddings": "nlp_roberta_sentence_embeddings",
+  "com.johnsnowlabs.nlp.embeddings.XlmRoBertaSentenceEmbeddings": "nlp_xlm_roberta_sentence_embeddings",
   "com.johnsnowlabs.nlp.embeddings.ChunkEmbeddings": "nlp_chunk_embeddings",
   "com.johnsnowlabs.nlp.embeddings.DistilBertEmbeddings": "nlp_distilbert_embeddings",
   "com.johnsnowlabs.nlp.embeddings.ElmoEmbeddings": "nlp_elmo_embeddings",

man/nlp_roberta_sentence_embeddings_pretrained.Rd

@@ -0,0 +1,30 @@
+setup({
+  sc <- testthat_spark_connection()
+  text_tbl <- testthat_tbl("test_text")
+
+  # These lines should set a pipeline that will ultimately create the columns needed for testing the annotator
+  assembler <- nlp_document_assembler(sc, input_col = "text", output_col = "document")
+
+  pipeline <- ml_pipeline(assembler)
+  test_data <- ml_fit_and_transform(pipeline, text_tbl)
+
+  assign("sc", sc, envir = parent.frame())
+  assign("pipeline", pipeline, envir = parent.frame())
+  assign("test_data", test_data, envir = parent.frame())
+})
+
+teardown({
+  spark_disconnect(sc)
+  rm(sc, envir = .GlobalEnv)
+  rm(pipeline, envir = .GlobalEnv)
+  rm(test_data, envir = .GlobalEnv)
+})
+
+test_that("nlp_roberta_sentence_embeddings pretrained", {
+  model <- nlp_roberta_sentence_embeddings_pretrained(sc, input_cols = c("document"), output_col = "roberta_sentence_embeddings")
+  transformed_data <- ml_transform(model, test_data)
+  expect_true("roberta_sentence_embeddings" %in% colnames(transformed_data))
+  
+  expect_true(inherits(model, "nlp_roberta_sentence_embeddings"))
+})
+

man/nlp_xlm_roberta_sentence_embeddings_pretrained.Rd

@@ -0,0 +1,30 @@
+setup({
+  sc <- testthat_spark_connection()
+  text_tbl <- testthat_tbl("test_text")
+
+  # These lines should set a pipeline that will ultimately create the columns needed for testing the annotator
+  assembler <- nlp_document_assembler(sc, input_col = "text", output_col = "document")
+
+  pipeline <- ml_pipeline(assembler)
+  test_data <- ml_fit_and_transform(pipeline, text_tbl)
+
+  assign("sc", sc, envir = parent.frame())
+  assign("pipeline", pipeline, envir = parent.frame())
+  assign("test_data", test_data, envir = parent.frame())
+})
+
+teardown({
+  spark_disconnect(sc)
+  rm(sc, envir = .GlobalEnv)
+  rm(pipeline, envir = .GlobalEnv)
+  rm(test_data, envir = .GlobalEnv)
+})
+
+test_that("nlp_bert_sentence_embeddings pretrained", {
+  model <- nlp_bert_sentence_embeddings_pretrained(sc, input_cols = c("document"), output_col = "bert_sentence_embeddings")
+  transformed_data <- ml_transform(model, test_data)
+  expect_true("bert_sentence_embeddings" %in% colnames(transformed_data))
+  
+  expect_true(inherits(model, "nlp_bert_sentence_embeddings"))
+})
+