com.flipkart.fdp.ml.transformer.StringIndexerTransformer Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of models-info Show documentation

There is a newer version: 0.4.0

package com.flipkart.fdp.ml.transformer;

import com.flipkart.fdp.ml.modelinfo.StringIndexerModelInfo;

import java.util.Map;

/**
 * Transforms input/ predicts for a String Indexer model representation
 * captured by  {@link com.flipkart.fdp.ml.modelinfo.StringIndexerModelInfo}.
 */
public class StringIndexerTransformer implements Transformer {

    private final StringIndexerModelInfo modelInfo;

    public StringIndexerTransformer(final StringIndexerModelInfo modelInfo) {
        this.modelInfo = modelInfo;
    }

    public double predict(final String s) {
        final Double index = modelInfo.getLabelToIndex().get(s);
        if (null == index) {
            throw new RuntimeException("Unseen label :" + s);
        }
        return index;
    }

    @Override
    public void transform(Map input) {
        String inp = (String) input.get(modelInfo.getInputKeys().iterator().next());
        input.put(modelInfo.getOutputKey(), predict(inp));
    }
}