docspell/project/NerModelsPlugin.scala

package docspell.build

import sbt.{Def, _}
import sbt.Keys._

/** Take some files from dependencies and put them into the resources
  * of a local sbt project.
  *
  * The reason is that the stanford ner model files are very very
  * large: the jar file for the english models is about 1G and the jar
  * file for the german models is about 170M. But I only need one file
  * that is about 60M from each jar. So just for the sake to save 1GB
  * file size when packaging docspell, this ugly plugin exists….
  *
  * The jar files to filter must be added to the libraryDependencies
  * in config "NerModels".
  */
object NerModelsPlugin extends AutoPlugin {

  object autoImport {
    val NerModels = config("NerModels")

    val nerModelsFilter = settingKey[String => Boolean]("Which files to keep.")
    val nerModelsRunFilter = taskKey[Seq[File]]("Extract files from libraryDependencies")

  }

  import autoImport._

  def nerModelSettings: Seq[Setting[_]] = Seq(
    nerModelsFilter := (_ => false),
    nerModelsRunFilter := {
      filterArtifacts(streams.value.log
        , Classpaths.managedJars(NerModels, Set("jar", "zip"), update.value)
        , nerModelsFilter.value
        , (Compile/resourceManaged).value)
    },
    Compile / resourceGenerators += nerModelsRunFilter.taskValue
  )

  def nerClassifierSettings: Seq[Setting[_]] = Seq(
    libraryDependencies ++= Dependencies.stanfordNlpModels.map(_ % NerModels),
    nerModelsFilter := {
      name => nerModels.exists(name.endsWith)
    }
  )

  override def projectConfigurations: Seq[Configuration] =
    Seq(NerModels)

  override def projectSettings: Seq[Setting[_]] =
    nerModelSettings

  def filterArtifacts(logger: Logger, cp: Classpath, nameFilter: NameFilter, out: File): Seq[File] = {
    logger.info(s"NerModels: Filtering artifacts...")
    cp.files.flatMap(f => {
      IO.unzip(f, out, nameFilter)
    })
  }

  private val nerModels = List(
    "german.conll.germeval2014.hgc_175m_600.crf.ser.gz",
    "english.all.3class.distsim.crf.ser.gz"
  )
}