com.flipkart.fdp.ml.adapter.ProbabilityTransformModelInfoAdapter Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of adapters-1.6_2.10 Show documentation

There is a newer version: 0.4.0

package com.flipkart.fdp.ml.adapter;

import com.flipkart.fdp.ml.modelinfo.ProbabilityTransformModelInfo;
import com.flipkart.fdp.ml.ProbabilityTransformModel;
import org.apache.spark.sql.DataFrame;

import java.util.LinkedHashSet;
import java.util.Set;

/**
 * Created by shubhranshu.shekhar on 18/08/16.
 */
public class ProbabilityTransformModelInfoAdapter extends AbstractModelInfoAdapter {
    @Override
    public ProbabilityTransformModelInfo getModelInfo(final ProbabilityTransformModel from, DataFrame df) {
        ProbabilityTransformModelInfo modelInfo = new ProbabilityTransformModelInfo();

        modelInfo.setActualClickProportion(from.getActualClickProportion());
        modelInfo.setUnderSampledClickProportion(from.getUnderSampledClickProportion());
        modelInfo.setProbIndex(from.getProbIndex());

        Set inputKeys = new LinkedHashSet();
        inputKeys.add(from.getInputCol());
        modelInfo.setInputKeys(inputKeys);

        Set outputKeys = new LinkedHashSet();
        outputKeys.add(from.getOutputCol());
        modelInfo.setOutputKeys(outputKeys);
        return modelInfo;
    }

    @Override
    public Class getSource() {
        return ProbabilityTransformModel.class;
    }

    @Override
    public Class getTarget() {
        return ProbabilityTransformModelInfo.class;
    }
}