|
import os |
|
from collections import Counter |
|
from multiprocessing import get_context |
|
|
|
from .fairseq_dictionary import Dictionary as fairseq_Dictionary |
|
|
|
|
|
class Dictionary(fairseq_Dictionary): |
|
"""Dictionary inheritted from FairSeq""" |
|
|
|
@staticmethod |
|
def _add_transcripts_to_dictionary_single_worker( |
|
transcripts, eos_word, worker_id=0, num_workers=1 |
|
): |
|
counter = Counter() |
|
size = len(transcripts) |
|
chunk_size = size // num_workers |
|
offset = worker_id * chunk_size |
|
end = min(size + 1, offset + chunk_size) |
|
for line in transcripts[offset:end]: |
|
for word in line.split(): |
|
counter.update([word]) |
|
counter.update([eos_word]) |
|
return counter |
|
|
|
@staticmethod |
|
def add_transcripts_to_dictionary(transcripts, dict, num_workers): |
|
def merge_result(counter): |
|
for w, c in sorted(counter.items()): |
|
dict.add_symbol(w, c) |
|
|
|
if num_workers > 1: |
|
pool = get_context("spawn").Pool(processes=num_workers) |
|
results = [] |
|
for worker_id in range(num_workers): |
|
results.append( |
|
pool.apply_async( |
|
Dictionary._add_transcripts_to_dictionary_single_worker, |
|
(transcripts, dict.eos_word, worker_id, num_workers), |
|
) |
|
) |
|
pool.close() |
|
pool.join() |
|
for r in results: |
|
merge_result(r.get()) |
|
else: |
|
merge_result( |
|
Dictionary._add_transcripts_to_dictionary_single_worker( |
|
transcripts, dict.eos_word |
|
) |
|
) |
|
|
