• Home
• com.johnsnowlabs.nlp
• spark-nlp_2.11
• 1.5.1
• source code
• PretrainedPipeline.scala

×

Project download

Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $

You can buy this project and download/modify it how often you want.

com.johnsnowlabs.nlp.pretrained.pipelines.PretrainedPipeline.scala Maven / Gradle / Ivy

package com.johnsnowlabs.nlp.pretrained.pipelines

import com.johnsnowlabs.nlp.pretrained.ResourceDownloader
import com.johnsnowlabs.nlp.{Finisher, LightPipeline}
import org.apache.spark.ml.PipelineModel
import org.apache.spark.sql.DataFrame

abstract class PretrainedPipeline(downloadName: String, language: Option[String] = None, folder: String = ResourceDownloader.publicFolder) {

  lazy protected val modelCache: PipelineModel = ResourceDownloader
    .downloadPipeline(downloadName, language, folder)

  def annotate(dataset: DataFrame, inputColumn: String): DataFrame = {
    modelCache
      .transform(dataset.withColumnRenamed(inputColumn, "text"))
  }

  def annotate(target: String): Map[String, Seq[String]] = new LightPipeline(modelCache).annotate(target)

  def annotate(target: Array[String]): Array[Map[String, Seq[String]]] = new LightPipeline(modelCache).annotate(target)

  def pretrained(): PipelineModel = modelCache

}