elastic / eland

Python Client and Toolkit for DataFrames, Big Data, Machine Learning and ETL in Elasticsearch
https://eland.readthedocs.io
Apache License 2.0
628 stars 98 forks source link

Simplify model import code #618

Open jeffvestal opened 9 months ago

jeffvestal commented 9 months ago

Feature Request

As a Python user, I would like to simplify the code required to load an inference model from Hugging Face (or private source) into Elasticsearch.

Use Case

Today eland is the way we load inference models into Elasticsearch. Many users will only use eland to load models, making it as simple to use will help the end-user experience.

Currently, in python, it requires several lines of code eg:

hf_model_id='sentence-transformers/msmarco-MiniLM-L-12-v3'
tm = TransformerModel(hf_model_id, "text_embedding")
es_model_id = tm.elasticsearch_model_id()
tmp_path = "models"
Path(tmp_path).mkdir(parents=True, exist_ok=True)
model_path, config, vocab_path = tm.save(tmp_path)
ptm = PyTorchModel(es, es_model_id)
ptm.import_model(model_path=model_path, config_path=None, vocab_path=vocab_path, config=config) 

This could be abstracted into a simple function call such as eland.ml.pytorch.import_model and take in the minimum required parameters. Something like

def import_model(
    es: Elasticsearch,
    model_id: str,
    task_type: str = 'text_embedding',
    access_token: Union[str, bool] = False
) -> str:

where it would return the model_id as it is known in Elasticsearch so it can be used later in the code to call inference

cc: @joshdevins