com.expleague.ml.models.multilabel.ConflictThresholdMultiLabelModel Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of jmll

Various ML methods implemented by myself and my students

There is a newer version: 1.4.9

package com.expleague.ml.models.multilabel;

import com.expleague.commons.math.vectors.Vec;
import com.expleague.commons.math.vectors.VecTools;
import com.expleague.commons.math.vectors.impl.vectors.ArrayVec;
import com.expleague.ml.models.multiclass.MCModel;

/**
 * User: qdeee
 * Date: 22.03.15
 */
public class ConflictThresholdMultiLabelModel extends MultiLabelModel.Stub {
  private final MCModel intern;
  private final double threshold;
  private final boolean allZeroesClassEnabled;

  public ConflictThresholdMultiLabelModel(final MCModel intern, final double threshold, final boolean allZeroesClassEnabled) {
    this.intern = intern;
    this.threshold = threshold;
    this.allZeroesClassEnabled = allZeroesClassEnabled;
  }

  @Override
  public Vec predictLabels(final Vec x) {
    final Vec prediction = intern.probs(x);
    final int argMax = VecTools.argmax(prediction);
    VecTools.toBinary(prediction, threshold);
    if (allZeroesClassEnabled) {
      if (argMax == prediction.dim() - 1) {
        //all zeroes class
        return new ArrayVec(ydim());
      } else {
        return prediction.sub(0, ydim());
      }
    } else {
      return prediction;
    }
  }

  @Override
  public int xdim() {
    return intern.xdim();
  }

  @Override
  public int ydim() {
    return allZeroesClassEnabled ? intern.countClasses() - 1 : intern.countClasses();
  }
}