com.intel.analytics.zoo.examples.chatbot.Utils.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of analytics-zoo-bigdl_0.12.2-spark_2.2.1 Show documentation

Big Data AI platform for distributed TensorFlow and PyTorch on Apache Spark.

The newest version!

/*
 * Copyright 2018 Analytics Zoo Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.intel.analytics.zoo.examples.chatbot

import scopt.OptionParser

object Utils {

  case class TrainParams(
    dataFolder: String = "./",
    batchSize: Int = 32,
    learningRate: Double = 0.0001,
    vocabSize: Int = 8004,
    nEpochs: Int = 20,
    embedDim: Int = 1024)

  val trainParser = new OptionParser[TrainParams]("Analytics Zoo chatbot Example") {
    opt[String]('f', "dataFolder")
      .text("where you put the text data")
      .action((x, c) => c.copy(dataFolder = x))
      .required()

    opt[Int]('b', "batchSize")
      .text("batchSize of rnn")
      .action((x, c) => c.copy(batchSize = x))

    opt[Double]('r', "learningRate")
      .text("learning rate")
      .action((x, c) => c.copy(learningRate = x))

    opt[Int]('e', "nEpochs")
      .text("epoch numbers")
      .action((x, c) => c.copy(nEpochs = x))

    opt[Int]("embedDims")
      .text("embedding dimensions")
      .action((x, c) => c.copy(embedDim = x))
  }
}