Boosting the speed of a translation model Helsinki-NLP/opus-mt-en-ar

🤗Transformers
1

Hi,

I am using Helsinki-NLP/opus-mt-en-ar as a trasnlation model and this is my code below.

I have profiled the code and found that the following is the bottleneck of he code

output_tokens = model.generate(**batch)

Is there any way to accelerate the code?

I am using pytorch cpu for this as I don’t have any GPU.

import pandas as pd
import torch
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, DataCollatorWithPadding
from torch.utils.data import Dataset, DataLoader
from line_profiler import profile
from optimum.bettertransformer import BetterTransformer

@profile
def function(): 
        
    df = pd.read_feather("data/ccs_synthetic.feather")

    # We will translate 1/6 of the dataset

    # 2092750

    df = df[:10]
    # df = df[:10]

    # from transformers import MarianTokenizer
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print("device: ", device)


    class CaptionDataset(Dataset):
        def __init__(self, df, tokenizer_name):
            self.df = df
            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)

        def __len__(self):
            return len(self.df)

        def __getitem__(self, index):
            sentence1 = df.loc[index, "caption"]

            tokens = self.tokenizer(sentence1, return_tensors="pt")

            return tokens


    tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
    model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
    # model = BetterTransformer.transform(model)

    # from transformers import AutoModelForSequenceClassification
    # from optimum.bettertransformer import BetterTransformer
    # model_hf = AutoModelForSequenceClassification.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
    # model = BetterTransformer.transform(model_hf, keep_original_model=True)

    model.to(device)
    model.eval()


    def custom_collate_fn(data):
        """
        Data collator with padding.
        """
        tokens = [sample["input_ids"][0] for sample in data]
        attention_masks = [sample["attention_mask"][0] for sample in data]

        attention_masks = torch.nn.utils.rnn.pad_sequence(attention_masks, batch_first=True)
        padded_tokens = torch.nn.utils.rnn.pad_sequence(tokens, batch_first=True)

        batch = {"input_ids": padded_tokens, "attention_mask": attention_masks}
        return batch


    print("After slicing: ", len(df))

    test_data = CaptionDataset(df, "Helsinki-NLP/opus-mt-en-ar")
    test_dataloader = DataLoader(
        test_data,
        batch_size=128,
        shuffle=False,
        # num_workers=0,
        collate_fn=custom_collate_fn,
    )

    with torch.no_grad():
        decoded_tokens = []
        for i, batch in enumerate(tqdm(test_dataloader)):

            batch = {k: v.to(device) for k, v in batch.items()}
            output_tokens = model.generate(**batch)
            decoded_tokens += tokenizer.batch_decode(
                output_tokens.to("cpu"), skip_special_tokens=True
            )


    df["caption_ar"] = decoded_tokens
    df.to_feather("data/ccs_synthetic_ar.feather")

if __name__ == "__main__":
    function()