Hi @NandiniN , thanks for taking a look! I'm linking below a reprex notebook in two formats: .py for running as a databricks notebook on DBR LTS, and a .ipynb notebook for running as a jupyter notebook natively (though I haven't tested this format). Let me know if you have issues getting it running.

https://drive.google.com/drive/folders/1V5hMzGlP3-nxXQUc-g8Y2qhZ3hs40ENs?usp=sharing