Huggingface hub integration 5545

[wwymak]

This is an initial draft of what the wrapper could look like.

Some questions I have:

• is this a sensible approach? (I have overridden the from_pretrained and save_pretrained methods so we can have defaults for the model cards, and a slightly different way of initialising the PyG model as some of them require the data at init time)
• should the wrapper live somewhere else in the codebase?
• do we want to have a nicer template for PyG models or just use the default one? Are there fields we want to have prepopulated by default
• should the wrapper also contain methods to log the training/evaluation metrics by default?
  @rusty1s would be great if you comment :)

The screenshot is an example of what it currently looks like if you use the wrapper to upload to huggingface model hub

[codecov]

Codecov Report

Merging #5930 (d96149a) into master (c0b10c0) will increase coverage by 0.02%.
The diff coverage is 100.00%.

❗ Current head d96149a differs from pull request most recent head ad7f405. Consider uploading reports for the commit ad7f405 to get more accurate results

@@            Coverage Diff             @@
##           master    #5930      +/-   ##
==========================================
+ Coverage   87.16%   87.18%   +0.02%     
==========================================
  Files         418      419       +1     
  Lines       22789    22838      +49     
==========================================
+ Hits        19863    19912      +49     
  Misses       2926     2926              

Impacted Files	Coverage Δ
torch_geometric/nn/model_hub.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[rusty1s]

This looks great. Happy to fine-tune it in later PRs as well.

is this a sensible approach? (I have overridden the from_pretrained and save_pretrained methods so we can have defaults for the model cards, and a slightly different way of initialising the PyG model as some of them require the data at init time)

I think so. Currently, we are creating a new model while I would prefer something like SchNet(ModelHubMixin).

should the wrapper live somewhere else in the codebase?

I think it should live in torch_geometric.nn rather than torch_geometric.nn.models.

do we want to have a nicer template for PyG models or just use the default one? Are there fields we want to have prepopulated by default

I think the default one works for now. Happy to fine-tune it in later PRs if that is possible. Please feel free to get creative :)

should the wrapper also contain methods to log the training/evaluation metrics by default?

I am not sure. Let's go with a basic version first.

[wwymak]

Hi @rusty1s I have converted the integration to a mixin, and added an example -- is this more along the lines of what you are thinking of?

[rusty1s]

Thanks, this is super cool. Wondering how we can integrate a basic test for this. Any idea?

[wwymak]

Wondering how we can integrate a basic test for this. Any idea?

I am thinking of a simple version of https://github.com/huggingface/huggingface_hub/blob/main/tests/test_hubmixin.py ? Since ours is based on what they've done for the pytorch mixin?

[wsad1]

Thanks again @wwymak for the great work. Will merge after tests pass.