Splitting the code in several files could enhance readability. For example could one file be dedicated to the prediction-type classes