Support for Converting Sentence-Transformers to TFLite

[adityasahugit]

System Info

optimum-1.23.3

Who can help?

@michaelbenayoun

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction (minimal, reproducible, runnable)

I am trying to export the model sentence-transformers/all-MiniLM-L6-v2 to TensorFlow Lite (TFLite) using the optimum-cli tool.

Since Sentence-Transformers often use BERT and other Transformer-based models as their backbone, I expected the export process to be supported. However, I encountered the following error:

optimum-cli export tflite --model sentence-transformers/all-MiniLM-L6-v2  --task semantic-segmentation tflite/

Gives error while Converting to TFLite

Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/tflite/__main__.py", line 147, in <module>
    main()
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/tflite/__main__.py", line 60, in main
    model = TasksManager.get_model_from_task(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/tasks.py", line 2071, in get_model_from_task
    model_class = TasksManager.get_model_class_for_task(
  File "/usr/local/lib/python3.10/dist-packages/optimum/exporters/tasks.py", line 1388, in get_model_class_for_task
    tasks_to_model_loader = TasksManager._LIBRARY_TO_TF_TASKS_TO_MODEL_LOADER_MAP[library]
KeyError: 'sentence_transformers'
Traceback (most recent call last):
  File "/usr/local/bin/optimum-cli", line 8, in <module>
    sys.exit(main())
  File "/usr/local/lib/python3.10/dist-packages/optimum/commands/optimum_cli.py", line 208, in main
    service.run()
  File "/usr/local/lib/python3.10/dist-packages/optimum/commands/export/tflite.py", line 243, in run
    subprocess.run(full_command, shell=True, check=True)
  File "/usr/lib/python3.10/subprocess.py", line 526, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command 'python3 -m optimum.exporters.tflite --model sentence-transformers/all-MiniLM-L6-v2 --task semantic-segmentation tflite/' returned non-zero exit status 1.

Expected behavior

Given that Sentence-Transformers are based on Transformer architectures such as BERT, RoBERTa, and MiniLM and are often used for feature extraction or embedding tasks, the optimum-cli tool should support exporting these models to TFLite

[github-actions]

This issue has been marked as stale because it has been open for 30 days with no activity. This thread will be automatically closed in 5 days if no further activity occurs.