[Refactor] Modified code for final data format