update to include train & test files