experimaestro / experimaestro-ir

IR module for experimaestro
GNU General Public License v3.0
11 stars 3 forks source link

Accessing pre-trained sub-modules #14

Closed bpiwowar closed 1 year ago

bpiwowar commented 1 year ago

When a model is trained, one might want to access one of its submodules in another task, e.g.

model = output["validation"]["best"]
faiss_index = FaissIndexBuilder(encoder=model.encoder).submit()

In that case, the encoder is random... and this is quite unexpected.

Potential solutions:

  1. Add a mechanism in experimaestro that allows linking the sub-models to the main one - when initialized, the encoder will require the model to be loaded from the checkpoint. Caveats: as with any implicit process, this makes xpmir more complex and might have other side effects.

  2. Add a mechanism to have several tasks to be run (and in that case, the output will be just the path), instead of just one, i.e.

    load_model = LoadModel(model=model, checkpoint=output["validation"]["best"])
    
    index = FaissIndexBuilder(encoder=model.encoder).submit(pre_tasks=[load_model])

    Important caveat: it breaks the current HuggingFace uploading/downloading mechanism and we need to include load_model each time we want to use the pre-trained model (which can begin to be cumbersome when various models are combined since we need to track down this). For HF, the solution would be to have this when loading

    model, load_model = AutoModel.load_from_hf_hub("xpmir/SPLADE_DistilMSE")
  3. Third solution, do something a bit in between

    best_model = model_loader(model=model, checkpoint=output["validation"]["best"])
    index = FaissIndexBuilder(encoder=best_model.encoder).submit()

    where model_loader is a special experimaestro object (ConfigFactory) that does whatever is needed and returns the "encoder" submodel.

bpiwowar commented 1 year ago

Re-opening: the current solution is not optimal, so it is better to switch to solution 2 (Add a mechanism to have several tasks to be run)

bpiwowar commented 1 year ago

Fixed in bcaa1eff7683236d502e0a79d5919a51c2eba702 with experimaestro 0.28 that includes initialization tasks.