Text Classification Python Interview Coding Question