NaiveBayesTextModelMapper extends RichModelMapper { @Override public void loadModel(List<Row> modelRows) { modelData = new NaiveBayesTextModelDataConverter().load(modelRows); } NaiveBayesTextModelMapper(TableSchema modelSchema, TableSchema dataSchema, Params params); @Override void loadModel(List<Row> modelRows); public String[] colNames; public String vectorColName; }
@Test public void testPredictCol() throws Exception { TableSchema dataSchema = new TableSchema( new String[] {"vec"}, new TypeInformation<?>[] {Types.STRING} ); Params params = new Params() .set(NaiveBayesTextPredictParams.VECTOR_COL, "vec") .set(NaiveBayesTextPredictParams.PREDICTION_COL, "pred") .set(NaiveBayesTextPredictParams.RESERVED_COLS, new String[] {}); NaiveBayesTextModelMapper mapper = new NaiveBayesTextModelMapper(modelSchema, dataSchema, params); mapper.loadModel(model); assertEquals(mapper.map(Row.of("1.0, 1.0, 0.0, 1.0")).getField(0), 1); assertEquals(mapper.getOutputSchema(), new TableSchema(new String[] {"pred"}, new TypeInformation<?>[] {Types.INT})); }
@Test public void testPredictReservedCol() throws Exception { TableSchema dataSchema = new TableSchema( new String[] {"vec"}, new TypeInformation<?>[] {Types.STRING} ); Params params = new Params() .set(NaiveBayesTextPredictParams.VECTOR_COL, "vec") .set(NaiveBayesTextPredictParams.PREDICTION_COL, "pred"); NaiveBayesTextModelMapper mapper = new NaiveBayesTextModelMapper(modelSchema, dataSchema, params); mapper.loadModel(model); assertEquals(mapper.map(Row.of("1.0, 1.0, 0.0, 1.0")).getField(1), 1); assertEquals(mapper.getOutputSchema(), new TableSchema(new String[] {"vec", "pred"}, new TypeInformation<?>[] {Types.STRING, Types.INT})); }