Update code in readme