Fixed generator class and added input class