ai.idylnlp.nlp.recognizer.configuration.DeepLearningEntityRecognizerConfiguration Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of idylnlp-nlp-entity-recognizers-deeplearning Show documentation
The newest version!
/*******************************************************************************
 * Copyright 2018 Mountain Fog, Inc.
 *
 * 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 ai.idylnlp.nlp.recognizer.configuration;

import java.io.File;
import java.util.LinkedHashSet;
import java.util.Set;

import ai.idylnlp.model.manifest.SecondGenModelManifest;
import ai.idylnlp.model.nlp.ConfidenceFilter;
import ai.idylnlp.model.nlp.configuration.AbstractEntityRecognizerConfiguration;
import ai.idylnlp.nlp.filters.confidence.SimpleConfidenceFilter;

import ai.idylnlp.nlp.recognizer.DeepLearningEntityRecognizer;

/**
 * Configuration for a {@link DeepLearningEntityRecognizer}.
 *
 * @author Mountain Fog, Inc.
 *
 */
public class DeepLearningEntityRecognizerConfiguration extends AbstractEntityRecognizerConfiguration {

  private String entityModelDirectory;

  public static class Builder {

    private ConfidenceFilter confidenceFilter;
    private Set blacklistedModelIDs;

    public Builder withConfidenceFilter(ConfidenceFilter confidenceFilter) {
      this.confidenceFilter = confidenceFilter;
      return this;
    }

    public Builder withBlacklistedModelIDs(Set blacklistedModelIDs) {
      this.blacklistedModelIDs = blacklistedModelIDs;
      return this;
    }

    /**
     * Creates the configuration.
     * @param entityModelDirectory The full path to the models directory.
     * @return A configured {@link DeepLearningEntityRecognizerConfiguration}.
     */
    public DeepLearningEntityRecognizerConfiguration build(String entityModelDirectory) {

      if(!entityModelDirectory.endsWith(File.separator)) {
        entityModelDirectory = entityModelDirectory + File.separator;
      }

      if(confidenceFilter == null) {
        confidenceFilter = new SimpleConfidenceFilter();
      }

      if(blacklistedModelIDs == null) {
        blacklistedModelIDs = new LinkedHashSet();
      }

      return new DeepLearningEntityRecognizerConfiguration(entityModelDirectory, confidenceFilter, blacklistedModelIDs);

    }

  }

  private DeepLearningEntityRecognizerConfiguration(
      String entityModelDirectory,
      ConfidenceFilter confidenceFilter,
      Set blacklistedModelIDs) {

    super(blacklistedModelIDs);

    this.entityModelDirectory = entityModelDirectory;
    this.blacklistedModelIDs = blacklistedModelIDs;
    this.confidenceFilter = confidenceFilter;

  }

  public String getEntityModelDirectory() {
    return entityModelDirectory;
  }

}