Text Generation

import spark.implicits._
import com.johnsnowlabs.nlp.base.DocumentAssembler
import com.johnsnowlabs.nlp.annotators.seq2seq.T5Transformer
import org.apache.spark.ml.Pipeline

// Step 1: Assembling the document from the input text
// Converts the input 'text' column into a 'document' column, required for NLP tasks
val documentAssembler = new DocumentAssembler()
  .setInputCol("text")
  .setOutputCol("documents")

// Step 3: Loading a pre-trained BERT model for token classification
// Applies a pre-trained BERT model for Named Entity Recognition (NER) to classify tokens
// `T5Transformer.pretrained()` loads the model, and `setInputCols` defines the input columns
val t5 = T5Transformer.pretrained("t5_small")
  .setTask("summarize:")
  .setInputCols(Array("documents"))
  .setMaxOutputLength(200)
  .setOutputCol("summaries")

// Step 4: Defining the pipeline
// The pipeline stages are document assembler, tokenizer, and token classifier
val pipeline = new Pipeline().setStages(Array(documentAssembler, t5))

// Step 5: Creating a sample DataFrame
// Creates a DataFrame with a sample sentence that will be processed by the pipeline
val data = Seq(
  "Transfer learning, where a model is first pre-trained on a data-rich task before being fine-tuned on a " +
  "downstream task, has emerged as a powerful technique in natural language processing (NLP). The effectiveness" +
  " of transfer learning has given rise to a diversity of approaches, methodology, and practice. In this " +
  "paper, we explore the landscape of transfer learning techniques for NLP by introducing a unified framework " +
  "that converts all text-based language problems into a text-to-text format. Our systematic study compares " +
  "pre-training objectives, architectures, unlabeled data sets, transfer approaches, and other factors on dozens " +
  "of language understanding tasks. By combining the insights from our exploration with scale and our new " +
  "Colossal Clean Crawled Corpus, we achieve state-of-the-art results on many benchmarks covering " +
  "summarization, question answering, text classification, and more. To facilitate future work on transfer " +
  "learning for NLP, we release our data set, pre-trained models, and code."
).toDF("text")

// Step 6: Fitting the pipeline and transforming the data
// The pipeline is fitted on the input data, then it performs the transformation to generate token labels
val result = pipeline.fit(data).transform(data)

// Step 7: Showing the results
// Displays the 'label.result' column, which contains the Named Entity Recognition (NER) labels for each token
result.select("summaries.result").show(false)

+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|result                                                                                                                                                                          |
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|[transfer learning has emerged as a powerful technique in natural language processing (NLP) the effectiveness of transfer learning has given rise to a diversity of approaches, |
|methodologies, and practice .]                                                                                                                                                  |
+--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+