Multiprocessing/Multithreading for huggingface pipeline

Viewed 733

I am using the zero shot classification pipeline provided by huggingface. I am trying to perform multiprocessing to parallelize the question answering. This is what I have tried till now

from pathos.multiprocessing import ProcessingPool as Pool
import multiprocess.context as ctx
from functools import partial
ctx._force_start_method('spawn') 

os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CUDA_LAUNCH_BLOCKING"]="1"

classifier = pipeline("zero-shot-classification", model="typeform/distilbert-base-uncased-mnli", device=0)

def get_label_score(classifier, labels,sentence):
    aa = classifier(sentence, labels)
    return aa['labels'][0],aa['scores'][0]

if __name__ == '__main__':
    pool = Pool(8)
    results = pool.map(
      partial(get_label_score, classifier,tag_values),
      d_lst)
    pool.close()
    pool.terminate()
    pool.join()

This does not improve the timing versus a for loop and results in OOM for list size greater than 10K sentences

I have also tried using executor.map this also does not improve the timings.

Please assist on how I can speed up inference

0 Answers
Related