feat: add multilanguage support to text2vec and vec2text

[jasonrichdarmawan]

Why ?

Why do we need to implement this feature ? What is the use case ?

To train SAEs using the NLLB's mined text from the slone/nllb-200-10M-sample dataset

For example, if you do dataset.take(2), sample 1 and sample 2 may have different lang1 and lang2. The current inference pipeline does not support this

How ?

Document the technical decisions you made.

1. if source_lang or target_lang is not instance of str, then the argument is assumed as Sequence, and so create tokenizer for every sample
   If some parts are WIP, please explicit them here.
2. this will create tokenizer for every sample instead of 1 tokenizer per batch

Instead of creating a tokenizer encoder for every text, it would be slightly more efficient to cache it for each language.
if you don't mind, i will submit a pull request today. However, I might need more time for the cache for each language

For context:

1. I am not sure why do you do tokenizer.create_encoder every time predict function is called instead of cache it in the first place e.g. is it expensive to store the tokenizer?
2. pytorch_lightning library has num_workers argument and has set_up function in the LightningModule. So, I assume the best practice is tokenizing should be done outside inference and generally be done in CPU

With that being said, that's why I don't touch how the tokenizer works, to avoid causing unexpected issue

Test plan

How did you test your changes ?

# %%
# Test autoencoder

if args["verbose"]:
  print("Testing text autoencoder...")
  text1_list, lang1_list = zip(
    *[
      (ex['text1'], ex['lang1']) 
      for ex in dataset.take(2)
    ]
  )
  embeddings = text2vec_model.predict(
    input=text1_list,
    source_lang=lang1_list,
  )
  print("Text encoder output shape:", embeddings.shape)

  text2_list, lang2_list = zip(
    *[
      (ex['text2'], ex['lang2']) 
      for ex in dataset.take(2)
    ]
  )
  prediction = vec2text_model.predict(
    inputs=embeddings,
    target_lang=lang2_list,
  )
  print(
    tabulate(
      [
        ["text1"] + list(text1_list),
        ["lang1"] + list(lang1_list),
        ["Prediction"] + list(prediction),
        ["text2"] + list(text2_list),
        ["lang2"] + list(lang2_list),
      ]
    )
  )

pytest /workspace/ALGOVERSE/UJR/jason/SONAR/tests/unit_tests/huggingface_pipelines/text.py

Problem: the unit tests failed in test_embedding_to_text_process_batch, test_text_to_embedding_to_text_pipeline_single_sentence, and test_text_to_embedding_to_text_pipeline_single_sentence. However, these use HF classes i.e. not part of this PR

Include full command line to help other people reproduce if needed.

[meta-cla]

Hi @jasonrichdarmawan!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[avidale]

It might be faster if you create encoders only once per language (e.g. by caching them with lru_cache or something like that).
Ideally, you should benchmark this.

[jasonrichdarmawan]

with functools.lru_cache(max_size=32), batch_size 10, the decoder took longer, how about for now we don't use cache?

[avidale]

To avoid having two big separate branches of code, maybe we simply convert the case of single language code into the case of sequence in the very beginning, and then proceed with the same translation function?

[jasonrichdarmawan]

how about this?