Changed to work with PRobabilityModels