Use UPath in helios/model.py

[joshhvulcan]

No description provided.

[favyen2]

Looks good to me.

load_model_and_optim_state only accepts paths on local filesystem, so it may be better to pass the original string to that function (or pass train_module_dir.path with a check that it is in fact on local filesystem).

It looks like currently (even before this PR) we just print "could not find helios encoder at ..." if the folder doesn't exist, which seems bad since we might erroneously fine-tune model without a checkpoint, but we can fix that when we figure out a better way to handle passing the Helios pre-training config+checkpoint to fine-tuning.

[joshhvulcan]

@favyen2 Updated to include .path at your suggestion. I've been looking through load_model_and_optim_state() and it appears to support reading remote files so I didn't add a local file check.

[favyen2]

@favyen2 Updated to include .path at your suggestion. I've been looking through load_model_and_optim_state() and it appears to support reading remote files so I didn't add a local file check.

Hm OK I was looking at https://github.com/allenai/OLMo-core/blob/main/src/olmo_core/io.py#L40 but I guess it may still be able to handle some protocols. Sorry in that case maybe it's better to pass the UPath or str(path) lol; because x.path will just include the path component not the protocol.

[joshhvulcan]

Hm OK I was looking at https://github.com/allenai/OLMo-core/blob/main/src/olmo_core/io.py#L40 but I guess it may still be able to handle some protocols. Sorry in that case maybe it's better to pass the UPath or str(path) lol; because x.path will just include the path component not the protocol.

I think that just strips off the file://, but will leave gs:// in place. load_model_and_optim_state() uses RemoteFileSystemReader to read the file which eventually drops into this block in get_bytes_range() which operates on the gs protocol. Anyway, I'm just going to pass the str() version through since that seems the safest at this point.