周末花了两天时间从安装xgboost到用于文本预测，记录下，首先是把文本分词，去停顿词，计算tf-idf值，然后模型训练，模型保存，加载模型，模型预测:
训练模型代码：
package com.meituan.model.xgboost; import java.util.Collections; import java.util.Comparator; import java.util.HashMap; import java.util.ArrayList; import java.util.List; import java.util.Arrays; import java.util.Map; import java.util.Map.Entry; import java.io.File; import java.io.IOException; import org.apache.commons.io.FileUtils; import ml.dmlc.xgboost4j.java.Booster; import ml.dmlc.xgboost4j.java.DMatrix; import ml.dmlc.xgboost4j.java.XGBoost; import ml.dmlc.xgboost4j.java.XGBoostError; public class TrainXgboost { private static String path = "/Users/shuubiasahi/Documents/workspace/xgboost/demo/data/"; private static String trainString = "agaricus.txt.train"; private static String testString = "agaricus.txt.test"; public static void main(String[] args) throws XGBoostError, IOException { DMatrix trainMat = new DMatrix("file/train.txt"); DMatrix testMat = new DMatrix("file/test.txt"); // specify parameters Map<String, Object> params = new HashMap<String, Object>(); params.put("booster", "gbtree"); params.put("eta", 0.6); // 为了防止过拟合，更新过程中用到的收缩步长。在每次提升计算之后，算法会直接获得新特征的权重 params.put("max_depth", 22);// 树最大深度 params.put("silent", 0); // 为1的时候不会打印模型迭代的信息，为0可以看到打印的信息 params.put("lambda", 2.5);// 用于逻辑回归的时候L2正则选项 params.put("min_child_weight", 6); // params.put("nthread", 6); //不使用的话系统会默认得到最大的线程数目 params.put("objective", "binary:logistic");// 目标函数值 // specify watchList HashMap<String, DMatrix> watches = new HashMap<String, DMatrix>(); watches.put("train", trainMat); // watches.put("test", testMat); // train a booster int round = 100; Booster booster = XGBoost.train(trainMat, params, round, watches, null, null); // booster.saveModel("xgboost/xgboost.model"); Map<String, Integer> map = booster.getFeatureScore(null); List<Map.Entry<String, Integer>> list = new ArrayList<Map.Entry<String, Integer>>( map.entrySet()); Collections.sort(list, new Comparator<Map.Entry<String, Integer>>() { @Override public int compare(Entry<String, Integer> o1, Entry<String, Integer> o2) { double result = o1.getValue() - o2.getValue(); if (result > 0) { return -1; } else { return 1; } } }); FileUtils.writeLines(new File("xgboost/keyword.txt"), list); float[][] result = booster.predict(testMat); /* * for(int i=0;i<result.length;i++){ * * for(int j=0;j<result[i].length;j++){ * System.out.print(result[i][j]+"\t"); } System.out.println(); } */ System.out.println("length is:" + result.length); } }
模型训练比较简单，先看看模型预测写的代码，准备的两个方法，把文本转化为libsvm的形式，再转化DMatrix：
public static CSRSparseData getSparseData(String content,Map<String, Terms> termsmap){ if (StringUtils.isBlank(content)) { return null; } Map<String, Long> maps = ToAnalysis .parse(WordUtil.replaceAllSynonyms(TextUtil.fan2Jian(WordUtil .replaceAll(content.toLowerCase())))) .getTerms() .stream() .map(x -> x.getName()) .filter(x -> !WordUtil.isStopword(x) ) .collect(Collectors.groupingBy(p -> p, Collectors.counting())); if (maps == null || maps.size() == 0) { return null; } int sum = maps.values().stream() .reduce((result, element) -> result = result + element).get() .intValue(); Map<Integer, Double> treemap = new TreeMap<Integer, Double>(); for (Entry<String, Long> map : maps.entrySet()) { String key = map.getKey(); Terms keyword = termsmap.get(key); double tf = TFIDF.tf(map.getValue(), sum); if (keyword == null) { continue; } int id = keyword.getId(); double idf = 0; idf = TFIDF.idf(termsmap.get("documentTotal").getFreq(), keyword.getFreq()); double tfidf = TFIDF.tfidf(tf, idf); treemap.put(id, tfidf); } if (treemap.size() == 0) { return null; } CSRSparseData spData = new CSRSparseData(); List<Float> tlabels = new ArrayList<>(); List<Float> tdata = new ArrayList<>(); List<Long> theaders = new ArrayList<>(); List<Integer> tindex = new ArrayList<>(); theaders.add(0l); theaders.add((long) treemap.size()); for (Entry<Integer, Double> map : treemap.entrySet()) { BigDecimal b = new BigDecimal(Double.toString(map.getValue())); tdata.add(b.floatValue()); tindex.add(Integer.valueOf(map.getKey())); } spData.labels = ArrayUtils.toPrimitive(tlabels .toArray(new Float[tlabels.size()])); spData.data = ArrayUtils.toPrimitive(tdata.toArray(new Float[tdata .size()])); spData.colIndex = ArrayUtils.toPrimitive(tindex .toArray(new Integer[tindex.size()])); spData.rowHeaders = ArrayUtils.toPrimitive(theaders .toArray(new Long[theaders.size()])); return spData; } public static double getClassification(Booster booster,String content,Map<String, Terms> termsmap) throws XGBoostError{ CSRSparseData spData=getSparseData(content, termsmap); if(spData==null){ return 0.0; } DMatrix data = new DMatrix(spData.rowHeaders, spData.colIndex, spData.data, DMatrix.SparseType.CSR, 0); return booster.predict(data)[0][0]; }
用于文本预测：
package com.meituan.model.xgboost; import java.io.BufferedReader; import java.io.BufferedWriter; import java.io.FileInputStream; import java.io.FileOutputStream; import java.io.IOException; import java.io.InputStreamReader; import java.io.OutputStreamWriter; import java.util.Map; import ml.dmlc.xgboost4j.java.Booster; import ml.dmlc.xgboost4j.java.XGBoost; import ml.dmlc.xgboost4j.java.XGBoostError; import com.meituan.model.libsvm.DocumentTransForm; import com.meituan.model.libsvm.Terms; public class Prediction { public static void main(String[] args) throws XGBoostError, IOException { Map<String, Terms> termsmap = DocumentTransForm.readmap("file/model"); Booster booster = XGBoost.loadModel("xgboost/xgboost.model"); System.out.println(DataLoader.getClassification(booster,"我们在吃饭",termsmap)); test(termsmap, booster, "/Users/shuubiasahi/Desktop/测试文件.csv"); } public static void test(Map<String, Terms> termsmap, Booster booster, String path) throws IOException, XGBoostError { BufferedReader buffer = new BufferedReader(new InputStreamReader( new FileInputStream(path))); BufferedWriter bufferwrite = new BufferedWriter(new OutputStreamWriter( new FileOutputStream("xgboost/merge.txt"))); BufferedWriter bufferwriteresult = new BufferedWriter( new OutputStreamWriter(new FileOutputStream( "xgboost/result.txt"))); String label = null; String line = buffer.readLine(); while (line != null) { String[] lines = line.split("\t"); if ("1".equalsIgnoreCase(lines[0]) && "美食".equalsIgnoreCase(lines[2])) { label = "1"; } else { label = "0"; } String content = lines[3]; double p = DataLoader.getClassification(booster, content, termsmap); if (p > 0) { // if (p > 0.86) { // if (WordUtil.isNumberMain(content) // ) { // p = 0.001; // } // } bufferwrite.write(label + "," + p + "\n"); String prString = p > 0.5 ? "1" : "0"; if (!label.equals(prString)) bufferwriteresult.write(label + "\t" + prString + "\t" + p + "\t" + lines[3] + "\n"); } line = buffer.readLine(); } bufferwriteresult.close(); buffer.close(); bufferwrite.close(); } }