<commit_before>package hex.tree.xgboost;

import biz.k11i.xgboost.Predictor;
import biz.k11i.xgboost.gbm.Dart;
import biz.k11i.xgboost.gbm.GBLinear;
import biz.k11i.xgboost.gbm.GBTree;
import biz.k11i.xgboost.tree.RegTree;
import biz.k11i.xgboost.tree.RegTreeNode;
import hex.LinkFunction;
import hex.LinkFunctionFactory;
import hex.genmodel.utils.LinkFunctionType;
import water.codegen.CodeGenerator;
import water.codegen.CodeGeneratorPipeline;
import water.exceptions.JCodeSB;
import water.util.SBPrintStream;

import static hex.genmodel.algos.xgboost.XGBoostMojoModel.ObjectiveType;

public abstract class XGBoostPojoWriter {

    public static XGBoostPojoWriter make(
        Predictor p,
        String namePrefix,
        XGBoostOutput output,
        double defaultThreshold
    ) {
        if (p.getBooster() instanceof GBTree) {
            return new XGBoostPojoTreeWriter(p, namePrefix, output, defaultThreshold);
        } else {
            return new XGBoostPojoLinearWriter(p, namePrefix, output, defaultThreshold);
        }
    }

    protected final Predictor _p;
    protected final String _namePrefix;
    protected final XGBoostOutput _output;
    private final double _defaultThreshold;

    protected XGBoostPojoWriter(Predictor p, String namePrefix, XGBoostOutput output, double defaultThreshold) {
        _p = p;
        _namePrefix = namePrefix;
        _output = output;
        _defaultThreshold = defaultThreshold;
    }
    
    protected String getFeatureAccessor(int idx) {
        if (idx >= _output._catOffsets[_output._cats]) {
            int colIdx = idx - _output._catOffsets[_output._cats] + _output._cats;
            if (_output._sparse) {
                return "(data[" + colIdx + "] == 0 ? Double.NaN : data[" + colIdx + "])";
            } else {
                return "data[" + colIdx + "]";
            }
        } else {
            int colIdx = 0;
            while (idx >= _output._catOffsets[colIdx + 1]) colIdx++;
            int colValue = idx - _output._catOffsets[colIdx];
            return "(data[" + colIdx + "] == " + colValue + " ? 1 : " + (_output._sparse ? "NaN" : "0") + ")";
        }
    }

    private void renderPredTransformViaLinkFunction(LinkFunctionType type, SBPrintStream sb) {
        LinkFunction lf = LinkFunctionFactory.getLinkFunction(type);
        sb.ip("preds[0] = (float) ").p(lf.linkInvString("preds[0]")).p(";").nl();
    }

    private void renderMultiClassPredTransform(SBPrintStream sb) {
        sb.ip("double max = preds[0];").nl();
        sb.ip("for (int i = 1; i < preds.length-1; i++) max = Math.max(preds[i], max); ").nl();
        sb.ip("double sum = 0.0D;").nl();
        sb.ip("for (int i = 0; i < preds.length-1; i++) {").nl();
        sb.ip("  preds[i] = Math.exp(preds[i] - max);").nl();
        sb.ip("  sum += preds[i];").nl();
        sb.ip("}").nl();
        sb.ip("for (int i = 0; i < preds.length-1; i++) {").nl();
        sb.ip("  preds[i] /= (float) sum;").nl();
        sb.ip("}").nl();
    }

    private void renderPredTransform(SBPrintStream sb) {<commit_msg>Fix XGBoost pojo writer when prediction is NaN.<commit_after>package hex.tree.xgboost;

import biz.k11i.xgboost.Predictor;
import biz.k11i.xgboost.gbm.Dart;
import biz.k11i.xgboost.gbm.GBLinear;
import biz.k11i.xgboost.gbm.GBTree;
import biz.k11i.xgboost.tree.RegTree;
import biz.k11i.xgboost.tree.RegTreeNode;
import hex.LinkFunction;
import hex.LinkFunctionFactory;
import hex.genmodel.utils.LinkFunctionType;
import water.codegen.CodeGenerator;
import water.codegen.CodeGeneratorPipeline;
import water.exceptions.JCodeSB;
import water.util.SBPrintStream;

import static hex.genmodel.algos.xgboost.XGBoostMojoModel.ObjectiveType;

public abstract class XGBoostPojoWriter {

    public static XGBoostPojoWriter make(
        Predictor p,
        String namePrefix,
        XGBoostOutput output,
        double defaultThreshold
    ) {
        if (p.getBooster() instanceof GBTree) {
            return new XGBoostPojoTreeWriter(p, namePrefix, output, defaultThreshold);
        } else {
            return new XGBoostPojoLinearWriter(p, namePrefix, output, defaultThreshold);
        }
    }

    protected final Predictor _p;
    protected final String _namePrefix;
    protected final XGBoostOutput _output;
    private final double _defaultThreshold;

    protected XGBoostPojoWriter(Predictor p, String namePrefix, XGBoostOutput output, double defaultThreshold) {
        _p = p;
        _namePrefix = namePrefix;
        _output = output;
        _defaultThreshold = defaultThreshold;
    }
    
    protected String getFeatureAccessor(int idx) {
        if (idx >= _output._catOffsets[_output._cats]) {
            int colIdx = idx - _output._catOffsets[_output._cats] + _output._cats;
            if (_output._sparse) {
                return "(data[" + colIdx + "] == 0 ? Double.NaN : data[" + colIdx + "])";
            } else {
                return "data[" + colIdx + "]";
            }
        } else {
            int colIdx = 0;
            while (idx >= _output._catOffsets[colIdx + 1]) colIdx++;
            int colValue = idx - _output._catOffsets[colIdx];
            return "(data[" + colIdx + "] == " + colValue + " ? 1 : " + (_output._sparse ? "Float.NaN" : "0") + ")";
        }
    }

    private void renderPredTransformViaLinkFunction(LinkFunctionType type, SBPrintStream sb) {
        LinkFunction lf = LinkFunctionFactory.getLinkFunction(type);
        sb.ip("preds[0] = (float) ").p(lf.linkInvString("preds[0]")).p(";").nl();
    }

    private void renderMultiClassPredTransform(SBPrintStream sb) {
        sb.ip("double max = preds[0];").nl();
        sb.ip("for (int i = 1; i < preds.length-1; i++) max = Math.max(preds[i], max); ").nl();
        sb.ip("double sum = 0.0D;").nl();
        sb.ip("for (int i = 0; i < preds.length-1; i++) {").nl();
        sb.ip("  preds[i] = Math.exp(preds[i] - max);").nl();
        sb.ip("  sum += preds[i];").nl();
        sb.ip("}").nl();
        sb.ip("for (int i = 0; i < preds.length-1; i++) {").nl();
        sb.ip("  preds[i] /= (float) sum;").nl();
        sb.ip("}").nl();
    }

    private void renderPredTransform(SBPrintStream sb) {
{'fraction_non_alphanumeric': 0.1396702230843841, 'fraction_numerical': 0.009052699644358229, 'mean_word_length': 2.736111111111111, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7252354', 'n_tokens_mistral': 2081, 'n_tokens_neox': 2027, 'n_words': 425}