Build a Disease Prediction AI in Java

Train a disease classifier in Java using Tribuo with Decision Trees and Random Forests.

Introduction

30 Second Summary

When you visit a doctor with symptoms, they mentally cross-reference your complaints against known disease patterns to narrow down a diagnosis. Machine learning models do the same thing, but with mathematical precision across thousands of cases.

In this project, you will build a disease prediction system entirely in Java that trains a classification model on symptom data, evaluates its accuracy, and lets you interactively input symptoms to receive a predicted diagnosis.

What You'll Build

A command-line Disease Prediction AI that loads a symptom-disease dataset, trains two classification models, compares their performance, and provides an interactive symptom checker.

By the end of this project, you'll have:

  • A Maven project with Tribuo as the ML dependency, requiring no external apps or pretrained models.
  • Two trained models (Decision Tree and Random Forest) with printed evaluation metrics including precision, recall, and F1 scores.
  • An interactive console-based symptom checker that accepts user input and outputs a predicted disease with confidence scores.
  • Secret Mission: Add a confidence threshold that warns the user when the model is uncertain about its prediction.

Are there any prerequisites?

This project assumes intermediate Java knowledge and familiarity with Maven projects in IntelliJ IDEA. You should have JDK 21 and IntelliJ IDEA already installed on your machine.

Before We Start

Before We Start...

Set Up the Maven Project and Dataset

Every machine learning system starts with three foundations: a project structure, a library to handle the ML logic, and data to learn from.

In this step, you will create a Maven project in IntelliJ IDEA, add Tribuo as the machine learning library, and prepare a dataset of symptoms mapped to diseases.

In this step, get ready to:
  • Create a Maven project with JDK 21 and add the Tribuo dependency.
  • Create the disease-symptom CSV dataset.
  • Verify the project compiles with a basic main class.
Create a new Maven project and add Tribuo

You need a Maven project to manage dependencies and build your Java code. Tribuo is an open-source ML library by Oracle Labs that provides classification algorithms out of the box.

  • In IntelliJ IDEA, click File in the top menu bar.
  • Click New, then Project.
  • In the New Project wizard, make sure Java is selected on the left.
  • Enter disease-predictor in the Name field.
  • Select Maven as the Build system.
  • Select your JDK 21 from the JDK dropdown.
  • Uncheck Add sample code if it is selected.
  • Expand Advanced Settings at the bottom of the dialog.
  • Set GroupId to com.example.
  • Set ArtifactId to disease-predictor.
  • Click Create.

What did this create?

IntelliJ generates a Maven project with a pom.xml file and the standard directory structure: src/main/java for your code, src/main/resources for data files, and src/test/java for tests.

The pom.xml file is where you declare project settings and external libraries that Maven downloads automatically.

IntelliJ may take a moment to index the new project. Wait until the progress bar at the bottom disappears before continuing.

  • Open pom.xml in the Project tool window on the left sidebar.
  • Replace the entire contents of pom.xml with the following code:
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>com.example</groupId>
    <artifactId>disease-predictor</artifactId>
    <version>1.0-SNAPSHOT</version>

    <properties>
        <maven.compiler.source>21</maven.compiler.source>
        <maven.compiler.target>21</maven.compiler.target>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    </properties>

    <dependencies>
        <!-- Tribuo ML library for classification -->
        <dependency>
            <groupId>org.tribuo</groupId>
            <artifactId>tribuo-all</artifactId>
            <version>4.3.2</version>
            <type>pom</type>
        </dependency>
    </dependencies>
</project>

What does this code do?

  • The <properties> block sets Java 21 as both the source and target compiler version.
  • The <dependency> block adds Tribuo 4.3.2, Oracle's open-source ML library. The tribuo-all artifact bundles all classification, regression, and data-loading modules in one import.
  • The <type>pom</type> tells Maven this is a BOM (Bill of Materials) dependency that pulls in transitive dependencies.
  • Save pom.xml with Cmd+S (macOS) or Ctrl+S (Windows).
  • When IntelliJ shows a popup asking to load Maven changes, click Load Maven Changes (or click the floating Maven refresh icon that appears in the editor).
  • Wait for the progress bar at the bottom to finish. Maven downloads the Tribuo library and its dependencies.

Maven sync failing?

  • Check your internet connection. Maven needs to download the Tribuo JAR files from Maven Central.
  • Make sure you selected JDK 21 when creating the project. Tribuo 4.3.2 supports Java 8+, but this project targets 21.
  • If you see a red underline on the dependency, right-click pom.xml in the Project tool window, select Maven, then Reload Project.

Why can't IntelliJ resolve the Tribuo dependency?

Create the disease-symptom dataset

Your ML model needs training data. You will create a CSV file where each row is a patient record with binary symptoms (1 = present, 0 = absent) and a disease label.

  • In the Project tool window, expand src, then main.
  • Right-click the resources folder.
  • Select New, then File.
  • Enter disease_symptoms.csv as the filename and press Enter.
  • Paste the following dataset into the new file:
fever,cough,fatigue,shortness_of_breath,body_aches,loss_of_taste,sore_throat,headache,nausea,disease
1,1,1,1,1,1,1,1,0,COVID-19
1,1,1,0,1,0,1,1,0,Influenza
0,1,0,0,0,0,1,0,0,Common Cold
1,0,1,0,1,0,0,1,1,Dengue
0,0,1,0,0,0,0,1,0,Anemia
1,0,1,0,1,0,0,1,1,Malaria
0,1,0,1,0,0,0,0,0,Asthma
0,0,0,0,0,0,0,1,1,Migraine
1,1,1,1,1,0,0,0,0,Pneumonia
1,1,1,1,1,1,0,1,0,COVID-19
1,1,1,0,1,0,0,1,0,Influenza
0,1,0,0,0,0,1,0,0,Common Cold
1,0,1,0,1,0,0,1,1,Dengue
0,0,1,0,0,0,0,0,0,Anemia
1,0,1,0,1,0,0,0,1,Malaria
0,0,0,1,0,0,0,0,0,Asthma
0,0,0,0,0,0,0,1,1,Migraine
1,1,1,1,0,0,1,0,0,Pneumonia
1,1,1,1,1,1,1,0,0,COVID-19
1,1,0,0,1,0,1,1,0,Influenza
0,0,0,0,0,0,1,0,0,Common Cold
1,0,1,0,1,0,0,1,0,Dengue
0,0,1,0,0,0,0,1,0,Anemia
1,0,0,0,1,0,0,1,1,Malaria
0,1,0,1,0,0,0,0,0,Asthma
0,0,0,0,0,0,0,1,0,Migraine
1,1,1,1,1,0,0,0,0,Pneumonia
1,1,1,0,1,1,1,1,0,COVID-19
1,1,1,0,1,0,1,0,0,Influenza
0,1,0,0,0,0,1,1,0,Common Cold
1,0,1,0,0,0,0,1,1,Dengue
0,0,1,0,0,0,0,0,1,Anemia
1,0,1,0,1,0,0,0,0,Malaria
0,1,0,1,0,0,1,0,0,Asthma
0,0,0,0,0,0,0,1,1,Migraine
1,1,1,1,0,0,1,0,0,Pneumonia

What does this dataset contain?

Each row represents a patient with 9 binary symptom columns (fever, cough, fatigue, shortness of breath, body aches, loss of taste, sore throat, headache, nausea) and one disease label column.

The dataset covers 9 diseases with 36 total records. This is a small educational dataset. In production, you would use thousands of records collected from clinical data.

  • Save disease_symptoms.csv.
  • Verify the file appears under src/main/resources in the Project tool window.

Don't see the resources folder?

  • If src/main/resources does not exist, right-click src/main, select New, then Directory, and type resources.
  • Make sure you created the file inside resources, not inside java. CSV data files belong in resources, not in your source code folder.

Help me find or create the resources folder in my Maven project.

Verify the project compiles

Before writing any ML logic, you should confirm that the project structure and dependencies are working. A simple main class that prints a message is enough to verify compilation.

  • In the Project tool window, expand src/main/java.
  • Right-click the java folder and select New, then Package.
  • Enter com.example and press Enter.
  • Right-click the new com.example package and select New, then Java Class.
  • Enter DiseasePredictor as the class name and press Enter.
  • Replace the contents of the file with the following code:
package com.example;

public class DiseasePredictor {
    public static void main(String[] args) {
        System.out.println("=== Disease Prediction AI ===");
    }
}

What does this code do?

This is a minimal Java class with a main method. It prints a startup message to confirm everything compiles and runs correctly. You will expand this class with ML logic in the next steps.

  • Save DiseasePredictor.java.
  • Right-click inside the editor and select Run 'DiseasePredictor.main()'.
  • Check the Run tool window at the bottom. You should see === Disease Prediction AI === printed in the output.

Seeing a compilation error?

  • Make sure the package declaration at the top of the file says package com.example; and that the file lives inside the com/example folder under src/main/java.
  • If IntelliJ shows red errors on the Tribuo imports (from later steps), confirm the Maven sync completed. Look for the green checkmark in the Maven tool window.
  • If the run option does not appear, make sure IntelliJ recognizes src/main/java as a sources root. Right-click java, select Mark Directory as, then Sources Root.

My DiseasePredictor class won't compile or run. Help me fix it.

✔️ Awesome, I've got everything!

Great. Make sure you have saved all three files before moving on.

ⓧ I'd like to double check the full code

Here is each file as it should look at the end of this step.

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>com.example</groupId>
    <artifactId>disease-predictor</artifactId>
    <version>1.0-SNAPSHOT</version>

    <properties>
        <maven.compiler.source>21</maven.compiler.source>
        <maven.compiler.target>21</maven.compiler.target>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    </properties>

    <dependencies>
        <dependency>
            <groupId>org.tribuo</groupId>
            <artifactId>tribuo-all</artifactId>
            <version>4.3.2</version>
            <type>pom</type>
        </dependency>
    </dependencies>
</project>
package com.example;

public class DiseasePredictor {
    public static void main(String[] args) {
        System.out.println("=== Disease Prediction AI ===");
    }
}
fever,cough,fatigue,shortness_of_breath,body_aches,loss_of_taste,sore_throat,headache,nausea,disease
1,1,1,1,1,1,1,1,0,COVID-19
1,1,1,0,1,0,1,1,0,Influenza
0,1,0,0,0,0,1,0,0,Common Cold
1,0,1,0,1,0,0,1,1,Dengue
0,0,1,0,0,0,0,1,0,Anemia
1,0,1,0,1,0,0,1,1,Malaria
0,1,0,1,0,0,0,0,0,Asthma
0,0,0,0,0,0,0,1,1,Migraine
1,1,1,1,1,0,0,0,0,Pneumonia
1,1,1,1,1,1,0,1,0,COVID-19
1,1,1,0,1,0,0,1,0,Influenza
0,1,0,0,0,0,1,0,0,Common Cold
1,0,1,0,1,0,0,1,1,Dengue
0,0,1,0,0,0,0,0,0,Anemia
1,0,1,0,1,0,0,0,1,Malaria
0,0,0,1,0,0,0,0,0,Asthma
0,0,0,0,0,0,0,1,1,Migraine
1,1,1,1,0,0,1,0,0,Pneumonia
1,1,1,1,1,1,1,0,0,COVID-19
1,1,0,0,1,0,1,1,0,Influenza
0,0,0,0,0,0,1,0,0,Common Cold
1,0,1,0,1,0,0,1,0,Dengue
0,0,1,0,0,0,0,1,0,Anemia
1,0,0,0,1,0,0,1,1,Malaria
0,1,0,1,0,0,0,0,0,Asthma
0,0,0,0,0,0,0,1,0,Migraine
1,1,1,1,1,0,0,0,0,Pneumonia
1,1,1,0,1,1,1,1,0,COVID-19
1,1,1,0,1,0,1,0,0,Influenza
0,1,0,0,0,0,1,1,0,Common Cold
1,0,1,0,0,0,0,1,1,Dengue
0,0,1,0,0,0,0,0,1,Anemia
1,0,1,0,1,0,0,0,0,Malaria
0,1,0,1,0,0,1,0,0,Asthma
0,0,0,0,0,0,0,1,1,Migraine
1,1,1,1,0,0,1,0,0,Pneumonia

Your project skeleton is ready with a working build, a machine learning library, and real training data. Next up, you will load this dataset into Tribuo and split it into training and testing sets.

Load and Explore the Dataset

In Step 1, you set up a Maven project with Tribuo as your ML library and created a dataset of symptoms mapped to diseases. Now it is time to teach your program how to read that data.

Before any machine learning model can learn patterns, it needs structured data loaded into memory. In this step, you will use Tribuo's data loading tools to read your CSV file, then split the data into separate training and testing sets so the model can learn from one portion and be evaluated on another.

In this step, get ready to:
  • Load the CSV dataset using Tribuo's CSVLoader and LabelFactory.
  • Split the data into training and testing sets.
  • Print dataset statistics to verify everything loaded correctly.
Load the CSV dataset with CSVLoader

Tribuo's CSVLoader reads numerical CSV files and converts them into a format the ML library can work with. It needs two things: a LabelFactory that tells it you are doing a classification task, and a list of column headers so it knows which column is the target (the disease) and which are features (the symptoms).

  • In IntelliJ IDEA, open DiseasePredictor.java from the src/main/java/com/example/ folder.
  • Replace the entire contents of the file with the following code:
package com.example;

import org.tribuo.*;
import org.tribuo.classification.*;
import org.tribuo.data.csv.CSVLoader;
import org.tribuo.evaluation.TrainTestSplitter;

import java.io.IOException;
import java.nio.file.Paths;

public class DiseasePredictor {

    // Define all column names including the target "disease" column
    private static final String[] HEADERS = {
        "fever", "cough", "fatigue", "shortness_of_breath",
        "body_aches", "loss_of_taste", "sore_throat", "headache", "nausea", "disease"
    };

    public static void main(String[] args) throws IOException {
        System.out.println("=== Disease Prediction AI ===");
        System.out.println("Loading dataset...\n");

        // LabelFactory tells Tribuo this is a classification problem
        var labelFactory = new LabelFactory();
        var csvLoader = new CSVLoader<>(labelFactory);

        // Load the CSV, specifying which column holds the label
        var dataSource = csvLoader.loadDataSource(
            Paths.get("src/main/resources/disease_symptoms.csv"),
            "disease",
            HEADERS
        );
    }
}

What does this code do?

This code sets up the data loading pipeline for your ML model.

  • HEADERS defines all column names in your CSV. The last entry "disease" is the target column the model will learn to predict.
  • LabelFactory tells Tribuo you are building a classifier (predicting a category) rather than a regressor (predicting a number).
  • loadDataSource() reads the CSV file, treats each row as one training example, and marks the "disease" column as the label to predict.

Why CSVLoader for this dataset?

Tribuo's CSVLoader is designed specifically for numerical CSV files with a String response field. Since your symptom data is all binary (0 or 1) and the disease column is a String, this loader is the perfect fit.

Split the data into training and testing sets

A model that is evaluated on the same data it trained on gives misleadingly high accuracy. Splitting the data into separate training and testing portions lets you measure how well the model handles examples it has never seen before.

  • In DiseasePredictor.java, find the closing ); of the loadDataSource call.
  • Add the following code below it, before the closing braces:
        // Split data into training (70%) and testing (30%) sets
        var splitter = new TrainTestSplitter<>(dataSource, 0.7, 1L);
        var trainData = new MutableDataset<>(splitter.getTrain());
        var testData = new MutableDataset<>(splitter.getTest());

        // Print dataset statistics to verify the load
        System.out.println("Training data size: " + trainData.size());
        System.out.println("Testing data size: " + testData.size());
        System.out.println("Number of features: " + trainData.getFeatureMap().size());
        System.out.println("Disease classes: " + trainData.getOutputIDInfo().getDomain());
        System.out.println();

What does this code do?

  • TrainTestSplitter randomly divides your dataset. The 0.7 means 70% goes to training and 30% to testing. The 1L is a random seed that makes the split reproducible.
  • MutableDataset wraps the raw data into Tribuo's dataset object that models can train on.
  • The print statements confirm how many examples ended up in each set, how many features (symptoms) the model sees, and which disease classes exist.

What is a train/test split?

Think of it like studying for an exam. The training set is your textbook (what you learn from). The testing set is the exam itself (questions you have never seen). If you only tested yourself on questions from the textbook, you would not know if you actually understood the material.

The 70/30 ratio is a common starting point in machine learning. It gives the model enough data to learn patterns while reserving enough to produce a meaningful evaluation.

  • Save DiseasePredictor.java.

✔️ Awesome, I've got everything!

Great. Double check you have saved your file before running it.

ⓧ I'd like to double check the full code

Here is the complete DiseasePredictor.java file at this point:

package com.example;

import org.tribuo.*;
import org.tribuo.classification.*;
import org.tribuo.data.csv.CSVLoader;
import org.tribuo.evaluation.TrainTestSplitter;

import java.io.IOException;
import java.nio.file.Paths;

public class DiseasePredictor {

    private static final String[] HEADERS = {
        "fever", "cough", "fatigue", "shortness_of_breath",
        "body_aches", "loss_of_taste", "sore_throat", "headache", "nausea", "disease"
    };

    public static void main(String[] args) throws IOException {
        System.out.println("=== Disease Prediction AI ===");
        System.out.println("Loading dataset...\n");

        var labelFactory = new LabelFactory();
        var csvLoader = new CSVLoader<>(labelFactory);

        var dataSource = csvLoader.loadDataSource(
            Paths.get("src/main/resources/disease_symptoms.csv"),
            "disease",
            HEADERS
        );

        var splitter = new TrainTestSplitter<>(dataSource, 0.7, 1L);
        var trainData = new MutableDataset<>(splitter.getTrain());
        var testData = new MutableDataset<>(splitter.getTest());

        System.out.println("Training data size: " + trainData.size());
        System.out.println("Testing data size: " + testData.size());
        System.out.println("Number of features: " + trainData.getFeatureMap().size());
        System.out.println("Disease classes: " + trainData.getOutputIDInfo().getDomain());
        System.out.println();
    }
}
Run the program and verify the dataset loads
  • In IntelliJ IDEA, click the green play button next to the main method, or press Ctrl+Shift+F10 (Windows) or Ctrl+Shift+R (macOS) to run the program.

You should see output similar to this in the Run panel:

=== Disease Prediction AI ===
Loading dataset...

Training data size: 25
Testing data size: 11
Number of features: 9
Disease classes: [Anemia, Asthma, COVID-19, Common Cold, Dengue, Influenza, Malaria, Migraine, Pneumonia]

What should I see?

  • Training data size: 25 means 70% of 36 rows (rounded) went to training.
  • Testing data size: 11 means the remaining 30% is reserved for evaluation.
  • Number of features: 9 confirms all 9 symptom columns are recognized as input features.
  • Disease classes lists all 9 unique diseases the model will learn to predict.

Seeing a FileNotFoundException?

  • Make sure disease_symptoms.csv is located at src/main/resources/disease_symptoms.csv. The path in the code must match exactly.
  • Confirm IntelliJ is running the program from the project root directory. Check your run configuration's working directory setting.
  • If you see a compilation error about IOException, make sure your main method signature includes throws IOException.

Help me fix my CSVLoader FileNotFoundException in Tribuo

Your data pipeline is working. The program can now load symptoms and diseases from a CSV file and split them into training and testing sets. Next up, you will train two different ML models on this data and compare which one predicts diseases more accurately.

Train and Compare Models

Your data is loaded and split into training and testing sets. Now it's time for the exciting part: training machine learning models that can learn the patterns between symptoms and diseases.

You'll train two different classification algorithms and compare their performance. A Decision Tree learns by asking yes/no questions about features. A Random Forest combines many decision trees to produce more robust predictions.

In this step, get ready to:
  • Train a Decision Tree classifier on the symptom data.
  • Train a Random Forest classifier with 10 ensemble trees.
  • Evaluate both models and compare their accuracy, precision, recall, and F1 scores.
Train a Decision Tree classifier

A Decision Tree works by learning a sequence of feature-based splits from your training data. Each node in the tree asks a question about one symptom, and the path from root to leaf produces a disease classification.

  • In DiseasePredictor.java, find the import section at the top of the file.
  • Add these new imports below the existing import org.tribuo.classification.*; line:
import org.tribuo.classification.dtree.CARTClassificationTrainer;
import org.tribuo.classification.evaluation.*;
import org.tribuo.common.tree.RandomForestTrainer;
import org.tribuo.ensemble.VotingCombiner;

What are these imports for?

Each import brings in a specific piece of Tribuo's classification toolkit.

  • CARTClassificationTrainer trains a single Decision Tree using the CART algorithm.
  • classification.evaluation.* provides the evaluator and evaluation results for measuring model accuracy.
  • RandomForestTrainer trains an ensemble of decision trees that vote together.
  • VotingCombiner tells the ensemble how to combine individual tree predictions into one answer.
  • Next, add a SYMPTOMS constant above the existing HEADERS array. Copy this code directly above the private static final String[] HEADERS line:
    private static final String[] SYMPTOMS = {
        "fever", "cough", "fatigue", "shortness_of_breath",
        "body_aches", "loss_of_taste", "sore_throat", "headache", "nausea"
    };

Why a separate SYMPTOMS array?

The HEADERS array includes the target column disease because the CSV loader needs it. The SYMPTOMS array lists only the input features, which you'll use later when asking the user about their symptoms.

  • Now add the Decision Tree training code. Find the last System.out.println(); line at the end of your main method (after the dataset statistics section).
  • Copy this code below that line:
        // Step 3: Train and compare models
        System.out.println("--- Training Decision Tree ---");
        var cartTrainer = new CARTClassificationTrainer();
        Model<Label> treeModel = cartTrainer.train(trainData);
        System.out.println("Decision Tree trained.\n");

What does this code do?

  • CARTClassificationTrainer() creates a trainer using the CART algorithm with default settings. It uses the Gini index to measure impurity at each split.
  • cartTrainer.train(trainData) feeds the training dataset into the algorithm. The trainer learns which symptom splits best separate the disease classes.
  • The result is a trained Model<Label> that can predict disease labels from new symptom inputs.
  • Save DiseasePredictor.java.
  • Run the program by clicking the green play button next to the main method. You should see the dataset statistics from before, followed by --- Training Decision Tree --- and Decision Tree trained. in the console.

Seeing a compilation error?

  • Check that all four new import lines are placed between the existing import org.tribuo.classification.*; and import org.tribuo.data.csv.CSVLoader; lines.
  • Make sure Model<Label> uses a capital M. Java generics are case-sensitive.
  • If IntelliJ shows red underlines on the imports, wait a moment for Maven to resolve the dependencies. You can also try clicking File then Invalidate Caches / Restart.

Help me fix my Decision Tree compilation error

Train a Random Forest classifier

A single Decision Tree can overfit to the training data. A Random Forest solves this by training many trees, each looking at a random subset of features, then combining their votes for a final prediction.

  • Add the Random Forest training code directly below the System.out.println("Decision Tree trained.\n"); line:
        System.out.println("--- Training Random Forest ---");
        var rfTreeTrainer = new CARTClassificationTrainer(6, 0.5f, 1L);
        var rfTrainer = new RandomForestTrainer<>(rfTreeTrainer, new VotingCombiner(), 10);
        Model<Label> forestModel = rfTrainer.train(trainData);
        System.out.println("Random Forest trained (10 trees).\n");

What does this code do?

  • CARTClassificationTrainer(6, 0.5f, 1L) creates a tree trainer with max depth 6, using only 50% of features at each split. The feature subsampling forces each tree to learn different patterns.
  • RandomForestTrainer<>(rfTreeTrainer, new VotingCombiner(), 10) builds an ensemble of 10 trees. Each tree trains on a random view of the data.
  • VotingCombiner uses majority voting across all 10 trees to pick the final disease prediction.

Why Random Forest over a single tree?

A single Decision Tree memorizes exact patterns in the training data. When it encounters slightly different symptom combinations at prediction time, it may get confused. A Random Forest averages out individual tree mistakes because each tree sees different feature subsets.

This is why the evaluation step next will likely show the Random Forest outperforming the single tree.

  • Save DiseasePredictor.java and run the program.
  • You should see both Decision Tree trained. and Random Forest trained (10 trees). printed in the console.

Seeing a PropertyException?

Tribuo's RandomForestTrainer requires the inner tree trainer to use feature subsampling. Make sure you're using new CARTClassificationTrainer(6, 0.5f, 1L) for rfTreeTrainer and not the default no-arg constructor.

Help me fix my Random Forest training error

Evaluate and compare both models

Training a model is only half the job. You need to measure how well each model performs on data it has never seen before. This is why you split the dataset earlier: the test set acts as unseen patient data.

  • Add the evaluation code below the System.out.println("Random Forest trained (10 trees).\n"); line:
        // Evaluate both models
        var evaluator = new LabelEvaluator();

        System.out.println("=== Decision Tree Evaluation ===");
        LabelEvaluation treeEval = evaluator.evaluate(treeModel, testData);
        System.out.println(treeEval.toString());
        System.out.println();

        System.out.println("=== Random Forest Evaluation ===");
        LabelEvaluation forestEval = evaluator.evaluate(forestModel, testData);
        System.out.println(forestEval.toString());
        System.out.println();

        // Pick the best model
        Model<Label> bestModel = forestEval.accuracy() >= treeEval.accuracy() ? forestModel : treeModel;
        String bestName = bestModel == forestModel ? "Random Forest" : "Decision Tree";
        System.out.println("Best model: " + bestName + " (accuracy: " +
            String.format("%.3f", Math.max(forestEval.accuracy(), treeEval.accuracy())) + ")");
        System.out.println();

What does this code do?

  • LabelEvaluator runs each model against the test set and calculates classification metrics.
  • evaluator.evaluate(model, testData) produces a LabelEvaluation containing accuracy, precision, recall, and F1 for every disease class.
  • The comparison logic picks whichever model has higher accuracy and stores it as bestModel. You'll use this best model for predictions in the next step.

What are precision, recall, and F1?

  • Precision answers: of all the times the model predicted a disease, how often was it correct?
  • Recall answers: of all the actual cases of a disease, how many did the model catch?
  • F1 is the harmonic mean of precision and recall. It gives you a single number that balances both.
  • Save DiseasePredictor.java and run the program.
  • You should see full evaluation tables for both models showing per-class precision, recall, and F1 scores. At the bottom, you should see a line starting with Best model: followed by the winning model name and its accuracy.

Seeing all zeros in the evaluation?

  • Make sure you're passing testData (not trainData) to the evaluate() method. Evaluating on training data gives misleading results.
  • Check that the CSV file still has all 36 data rows. If the file was accidentally modified, the splitter may produce an empty test set.

Help me understand my model evaluation results

✔️ Awesome, I've got everything!

Great. Double check you've saved your file and that both evaluation reports display correctly in the console.

ⓧ I'd like to double check the full code

Here's the complete DiseasePredictor.java file at this point:

package com.example;

import org.tribuo.*;
import org.tribuo.classification.*;
import org.tribuo.classification.dtree.CARTClassificationTrainer;
import org.tribuo.classification.evaluation.*;
import org.tribuo.common.tree.RandomForestTrainer;
import org.tribuo.data.csv.CSVLoader;
import org.tribuo.ensemble.VotingCombiner;
import org.tribuo.evaluation.TrainTestSplitter;

import java.io.IOException;
import java.nio.file.Paths;

public class DiseasePredictor {

    private static final String[] SYMPTOMS = {
        "fever", "cough", "fatigue", "shortness_of_breath",
        "body_aches", "loss_of_taste", "sore_throat", "headache", "nausea"
    };

    private static final String[] HEADERS = {
        "fever", "cough", "fatigue", "shortness_of_breath",
        "body_aches", "loss_of_taste", "sore_throat", "headache", "nausea", "disease"
    };

    public static void main(String[] args) throws IOException {
        System.out.println("=== Disease Prediction AI ===");
        System.out.println("Loading dataset...\n");

        var labelFactory = new LabelFactory();
        var csvLoader = new CSVLoader<>(labelFactory);

        var dataSource = csvLoader.loadDataSource(
            Paths.get("src/main/resources/disease_symptoms.csv"),
            "disease",
            HEADERS
        );

        var splitter = new TrainTestSplitter<>(dataSource, 0.7, 1L);
        var trainData = new MutableDataset<>(splitter.getTrain());
        var testData = new MutableDataset<>(splitter.getTest());

        System.out.println("Training data size: " + trainData.size());
        System.out.println("Testing data size: " + testData.size());
        System.out.println("Number of features: " + trainData.getFeatureMap().size());
        System.out.println("Disease classes: " + trainData.getOutputIDInfo().getDomain());
        System.out.println();

        System.out.println("--- Training Decision Tree ---");
        var cartTrainer = new CARTClassificationTrainer();
        Model<Label> treeModel = cartTrainer.train(trainData);
        System.out.println("Decision Tree trained.\n");

        System.out.println("--- Training Random Forest ---");
        var rfTreeTrainer = new CARTClassificationTrainer(6, 0.5f, 1L);
        var rfTrainer = new RandomForestTrainer<>(rfTreeTrainer, new VotingCombiner(), 10);
        Model<Label> forestModel = rfTrainer.train(trainData);
        System.out.println("Random Forest trained (10 trees).\n");

        var evaluator = new LabelEvaluator();

        System.out.println("=== Decision Tree Evaluation ===");
        LabelEvaluation treeEval = evaluator.evaluate(treeModel, testData);
        System.out.println(treeEval.toString());
        System.out.println();

        System.out.println("=== Random Forest Evaluation ===");
        LabelEvaluation forestEval = evaluator.evaluate(forestModel, testData);
        System.out.println(forestEval.toString());
        System.out.println();

        Model<Label> bestModel = forestEval.accuracy() >= treeEval.accuracy() ? forestModel : treeModel;
        String bestName = bestModel == forestModel ? "Random Forest" : "Decision Tree";
        System.out.println("Best model: " + bestName + " (accuracy: " +
            String.format("%.3f", Math.max(forestEval.accuracy(), treeEval.accuracy())) + ")");
        System.out.println();
    }
}

You now have two trained models and a clear winner based on accuracy. Next up, you'll build an interactive symptom checker that uses your best model to predict diseases from real-time user input.

Build the Interactive Symptom Checker

Your models are trained, evaluated, and the best performer is selected. The evaluation metrics prove the system can classify diseases from symptoms.

Now you'll put that trained model to work by building a real-time console interface. The user will answer yes/no for each symptom, and the model will return a predicted disease with confidence scores.

In this step, get ready to:
  • Create a console input loop that collects symptom responses.
  • Feed the collected symptoms into the trained model for prediction.
  • Display the predicted disease and probability distribution.
Add the new imports

The interactive checker needs two new classes from Java's standard library. Scanner reads console input, and Map holds the probability distribution across all diseases.

  • In DiseasePredictor.java, find your existing import block at the top of the file.
  • Add these two import lines after import java.nio.file.Paths;:
import java.util.Map;
import java.util.Scanner;

What do these imports do?

Scanner lets your program read text input from the terminal line by line. Map is needed to work with the key-value pairs of disease names and their predicted probabilities.

Create the symptom input loop

The interactive checker needs to ask the user about each symptom, convert their answer into a numerical feature value, and package those values into a format the model understands.

  • In DiseasePredictor.java, find the last System.out.println(); at the bottom of your main method (after the best model print statement).
  • Add the following code below that line to create the input loop:
        // Step 4: Interactive symptom checker
        System.out.println("=== Interactive Symptom Checker ===");
        System.out.println("Answer yes/no for each symptom.\n");

        Scanner scanner = new Scanner(System.in);
        boolean running = true;

        while (running) {
            var example = new ArrayExample<>(new Label("UNKNOWN"));

            for (String symptom : SYMPTOMS) {
                System.out.print("Do you have " + symptom.replace("_", " ") + "? (yes/no): ");
                String answer = scanner.nextLine().trim().toLowerCase();
                double value = answer.equals("yes") || answer.equals("y") ? 1.0 : 0.0;
                example.add(symptom, value);
            }

What does this code do?

  • ArrayExample is Tribuo's way of building a single data point programmatically. The placeholder Label("UNKNOWN") is required by the constructor but will be replaced by the model's prediction.
  • The for-loop iterates through each symptom name, prints a human-readable prompt (replacing underscores with spaces), and converts the user's yes/no answer into 1.0 or 0.0.
  • example.add(symptom, value) attaches each feature to the example using the same names the model was trained on. This is how Tribuo matches user input to the trained feature map.
Pass symptoms to the model for prediction

With the symptom features collected into an ArrayExample, you can now pass it to the best model's predict() method to get a disease classification.

  • Add this code directly after the closing brace of the for loop (still inside the while loop):
            Prediction<Label> prediction = bestModel.predict(example);
            Label predicted = prediction.getOutput();

What does this code do?

bestModel.predict(example) runs the trained model on your new symptom data and returns a Prediction object. This object contains both the top predicted label and the full probability distribution across all disease classes.

getOutput() extracts the most likely disease label from the prediction. The label carries both the disease name and its confidence score.

Display the prediction results

The final piece shows the user their predicted disease, the model's confidence, and the probability distribution across all possible diseases.

  • Add this code directly after the Label predicted = prediction.getOutput(); line:
            System.out.println("\n--- Prediction Result ---");
            System.out.println("Predicted disease: " + predicted.getLabel());
            System.out.println("Confidence: " + String.format("%.1f%%", predicted.getScore() * 100));

            System.out.println("\nAll probabilities:");
            Map<String, Label> outputScores = prediction.getOutputScores();
            outputScores.entrySet().stream()
                .sorted((a, b) -> Double.compare(b.getValue().getScore(), a.getValue().getScore()))
                .forEach(entry -> System.out.println("  " + entry.getKey() + ": " +
                    String.format("%.1f%%", entry.getValue().getScore() * 100)));

            System.out.print("\nCheck another set of symptoms? (yes/no): ");
            String again = scanner.nextLine().trim().toLowerCase();
            running = again.equals("yes") || again.equals("y");
            System.out.println();
        }

        System.out.println("Goodbye!");
        scanner.close();

What does this code do?

  • predicted.getLabel() returns the disease name as a string, and predicted.getScore() returns the confidence as a value between 0 and 1 (multiplied by 100 for a percentage).
  • getOutputScores() returns a map of every disease class and its probability. The stream sorts them in descending order so the most likely diseases appear first.
  • The continue prompt lets the user check multiple symptom sets without restarting the program. Setting running = false exits the while loop and triggers the cleanup.
  • Save DiseasePredictor.java.

Seeing a compile error?

  • If you see an error about ArrayExample not being found, make sure your wildcard import import org.tribuo.*; is still at the top of the file.
  • If Map or Scanner shows red, confirm you added both new imports after import java.nio.file.Paths;.
  • Check that every opening brace has a matching closing brace. The while loop and the for loop each need their own closing }.

Help me fix a compile error in my DiseasePredictor interactive loop

  • Run the program by clicking the green play button next to public static void main in IntelliJ.

The program will load the dataset, train both models, print evaluations, then drop you into the interactive symptom checker. Answer each symptom prompt with yes or no.

  • Try entering yes for fever, cough, fatigue, shortness of breath, and body aches. Enter no for the rest.
  • Verify you see a predicted disease name, a confidence percentage, and a list of all disease probabilities sorted from highest to lowest.
  • Enter no at the "Check another set of symptoms?" prompt to exit.

Program hangs or no input prompt appears?

  • Make sure you're clicking inside the IntelliJ **Run** console panel at the bottom of the screen before typing. The program waits for input there, not in the editor.
  • If you see the evaluation output but no "Do you have fever?" prompt, scroll down in the console. The training output may push the prompt out of view.

Help me troubleshoot my interactive symptom checker not accepting input

✔️ Awesome, I've got everything!

Great. Double check you've saved your file before moving on.

ⓧ I'd like to double check the full code

package com.example;

import org.tribuo.*;
import org.tribuo.classification.*;
import org.tribuo.classification.dtree.CARTClassificationTrainer;
import org.tribuo.classification.evaluation.*;
import org.tribuo.common.tree.RandomForestTrainer;
import org.tribuo.data.csv.CSVLoader;
import org.tribuo.ensemble.VotingCombiner;
import org.tribuo.evaluation.TrainTestSplitter;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.Map;
import java.util.Scanner;

public class DiseasePredictor {

    private static final String[] SYMPTOMS = {
        "fever", "cough", "fatigue", "shortness_of_breath",
        "body_aches", "loss_of_taste", "sore_throat", "headache", "nausea"
    };

    private static final String[] HEADERS = {
        "fever", "cough", "fatigue", "shortness_of_breath",
        "body_aches", "loss_of_taste", "sore_throat", "headache", "nausea", "disease"
    };

    public static void main(String[] args) throws IOException {
        System.out.println("=== Disease Prediction AI ===");
        System.out.println("Loading dataset...\n");

        var labelFactory = new LabelFactory();
        var csvLoader = new CSVLoader<>(labelFactory);

        var dataSource = csvLoader.loadDataSource(
            Paths.get("src/main/resources/disease_symptoms.csv"),
            "disease",
            HEADERS
        );

        var splitter = new TrainTestSplitter<>(dataSource, 0.7, 1L);
        var trainData = new MutableDataset<>(splitter.getTrain());
        var testData = new MutableDataset<>(splitter.getTest());

        System.out.println("Training data size: " + trainData.size());
        System.out.println("Testing data size: " + testData.size());
        System.out.println("Number of features: " + trainData.getFeatureMap().size());
        System.out.println("Disease classes: " + trainData.getOutputIDInfo().getDomain());
        System.out.println();

        System.out.println("--- Training Decision Tree ---");
        var cartTrainer = new CARTClassificationTrainer();
        Model<Label> treeModel = cartTrainer.train(trainData);
        System.out.println("Decision Tree trained.\n");

        System.out.println("--- Training Random Forest ---");
        var rfTreeTrainer = new CARTClassificationTrainer(6, 0.5f, 1L);
        var rfTrainer = new RandomForestTrainer<>(rfTreeTrainer, new VotingCombiner(), 10);
        Model<Label> forestModel = rfTrainer.train(trainData);
        System.out.println("Random Forest trained (10 trees).\n");

        var evaluator = new LabelEvaluator();

        System.out.println("=== Decision Tree Evaluation ===");
        LabelEvaluation treeEval = evaluator.evaluate(treeModel, testData);
        System.out.println(treeEval.toString());
        System.out.println();

        System.out.println("=== Random Forest Evaluation ===");
        LabelEvaluation forestEval = evaluator.evaluate(forestModel, testData);
        System.out.println(forestEval.toString());
        System.out.println();

        Model<Label> bestModel = forestEval.accuracy() >= treeEval.accuracy() ? forestModel : treeModel;
        String bestName = bestModel == forestModel ? "Random Forest" : "Decision Tree";
        System.out.println("Best model: " + bestName + " (accuracy: " +
            String.format("%.3f", Math.max(forestEval.accuracy(), treeEval.accuracy())) + ")");
        System.out.println();

        System.out.println("=== Interactive Symptom Checker ===");
        System.out.println("Answer yes/no for each symptom.\n");

        Scanner scanner = new Scanner(System.in);
        boolean running = true;

        while (running) {
            var example = new ArrayExample<>(new Label("UNKNOWN"));

            for (String symptom : SYMPTOMS) {
                System.out.print("Do you have " + symptom.replace("_", " ") + "? (yes/no): ");
                String answer = scanner.nextLine().trim().toLowerCase();
                double value = answer.equals("yes") || answer.equals("y") ? 1.0 : 0.0;
                example.add(symptom, value);
            }

            Prediction<Label> prediction = bestModel.predict(example);
            Label predicted = prediction.getOutput();

            System.out.println("\n--- Prediction Result ---");
            System.out.println("Predicted disease: " + predicted.getLabel());
            System.out.println("Confidence: " + String.format("%.1f%%", predicted.getScore() * 100));

            System.out.println("\nAll probabilities:");
            Map<String, Label> outputScores = prediction.getOutputScores();
            outputScores.entrySet().stream()
                .sorted((a, b) -> Double.compare(b.getValue().getScore(), a.getValue().getScore()))
                .forEach(entry -> System.out.println("  " + entry.getKey() + ": " +
                    String.format("%.1f%%", entry.getValue().getScore() * 100)));

            System.out.print("\nCheck another set of symptoms? (yes/no): ");
            String again = scanner.nextLine().trim().toLowerCase();
            running = again.equals("yes") || again.equals("y");
            System.out.println();
        }

        System.out.println("Goodbye!");
        scanner.close();
    }
}

Your disease predictor is now fully functional. It loads data, trains two models, picks the best one, and lets you interact with it in real time. Next up, you'll push this further in the Secret Mission by adding confidence thresholds that warn users when the model is uncertain.

Secret mission

Add Confidence Thresholds

Your predictor currently delivers every diagnosis with equal certainty. In the real world, a model that says "maybe Dengue, maybe Malaria" with 30% confidence each should behave very differently from one that says "COVID-19" with 90% confidence. Add a confidence threshold that warns users when the model is unsure.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing costs. Decide whether to keep your project files for future reference or delete them to free up disk space.

Resources you used:

  • The disease-predictor project folder (Java source, CSV dataset, Maven config).
  • Tribuo dependency jars cached in your local Maven repository (~/.m2/repository/org/tribuo/).

Keep everything running

No action needed. Choose this if you want to keep experimenting with different symptom combinations, tweak the model parameters, or expand the dataset.

  • Your project stays in IntelliJ IDEA and you can run it any time.
  • The cached Maven dependencies take up a small amount of disk space but have no ongoing cost.

Pause - I'll come back to this later

Close IntelliJ IDEA to free up memory. Your project files stay on disk and are ready to reopen whenever you return.

  • Close the disease-predictor project in IntelliJ IDEA.
  • To resume later, reopen IntelliJ IDEA and select the disease-predictor project from your recent projects list.

Delete - I don't want to use this again

Remove all project files and cached dependencies to start fresh.

  • Close the project in IntelliJ IDEA.
  • Delete the project folder by running this command (macOS/Linux):
rm -rf ~/IdeaProjects/disease-predictor

On Windows, run this instead:

rmdir /s /q %USERPROFILE%\IdeaProjects\disease-predictor
  • Optionally, remove the cached Tribuo jars from your local Maven repository by running (macOS/Linux):
rm -rf ~/.m2/repository/org/tribuo

On Windows:

rmdir /s /q %USERPROFILE%\.m2\repository\org\tribuo

What about IntelliJ IDEA and JDK 21?

You already had IntelliJ IDEA, JDK 21, and Maven installed before this project. Keep them for future Java development. They do not consume ongoing resources.

That's a Wrap!

That's a Wrap!

Nice work! You just built a fully functional disease prediction system in Java, training machine learning models from scratch using nothing but a CSV file and the Tribuo library.

You've learned how to:

  • Built a Maven project with Tribuo and trained two classification models (Decision Tree and Random Forest) from a symptom-disease CSV dataset.
  • Evaluated model performance using precision, recall, F1, and accuracy metrics to automatically select the best-performing model.
  • Built an interactive console-based symptom checker that accepts user input and outputs predicted diseases with probability distributions.
  • Secret Mission: Added confidence thresholds that warn users when the model is uncertain about its prediction.

Ready to quiz yourself?