Updated to generator v2 code