The plotting code can live here for now, let's use hugging face to pull the model and make sure it runs in colab.