How to add a progress bar during embedding

[rangehow]

I am trying following code to utilize multiple gpu, but add tqdm in outer seems not that good (because the progress bar have a total device_count which makes it hard to watch).
So if there any inner args that can open progress bar?

import asyncio
import os
from infinity_emb import AsyncEngineArray, EngineArgs, AsyncEmbeddingEngine
from torch.cuda import device_count
import numpy as np
import argparse
import datasets


# sentences = ["Paris is in France."]



# Split sentences into roughly equal chunks for each engine
def split_sentences(sentences, engine_count):
    # Split sentences into `engine_count` chunks
    avg_len = len(sentences) // engine_count
    chunks = [sentences[i * avg_len: (i + 1) * avg_len] for i in range(engine_count)]
    # If there are leftovers, add them to the last chunk
    if len(sentences) % engine_count != 0:
        chunks[-1].extend(sentences[engine_count * avg_len:])
    return chunks


async def embed_text(engine: AsyncEmbeddingEngine, sentences_chunk: list): 
    async with engine:
        embeddings, _ = await engine.embed(sentences=sentences_chunk)
    return embeddings

async def run_parallel_embeddings(sentences):
    # Split the sentences into chunks for each engine
    sentence_chunks = split_sentences(sentences, engine_count)

    tasks = []
    # Create a task for each engine, passing the respective chunk of sentences
    for i in range(engine_count):

        tasks.append(embed_text(array[i], sentence_chunks[i]))
        
    # Wait for all embeddings to complete
    results = await asyncio.gather(*tasks)

    combined_embeddings = np.concatenate([result for result in results], axis=0)
    
    return combined_embeddings




def parse_dataset(dataset_str):
    if dataset_str.endswith('.tsv'):
        # 尝试读取TSV文件
        dataset = datasets.load_dataset("csv", data_files=dataset_str, delimiter="\t", num_proc=32)['train']
    # elif dataset_str.endwith('.jsonl'):

    else:
        dataset = datasets.load_from_disk(dataset_str)
    return dataset


def parse_args():
    parser = argparse.ArgumentParser(description="Process datasets with padding and multiprocessing.")
    parser.add_argument("--model", help="Model name or path.")
    parser.add_argument("--dataset", required=True, help="Path to the dataset.")
    parser.add_argument("--batch_size", type=int, default=1024, help="Batch size per worker.")
    parser.add_argument("--output_dir", default="./embeddings", help="Directory to save embeddings.")
    parser.add_argument("--model_name")
    return parser.parse_args()

if __name__ == "__main__":
    
    engine_count = device_count()

    args = parse_args()


    # 获取倒数第二个部分
    model_name = args.model_name
    print(model_name)
    # 载入数据集
    dataset = parse_dataset(args.dataset)


    # Set up multiple engines (one for each GPU)
    array = AsyncEngineArray.from_args([
        EngineArgs(
            model_name_or_path=args.model,
            engine="torch",
            embedding_dtype="float32",
            dtype="auto",
            device='cuda',
            device_id=f'{i}',
            served_model_name=f'model_{i}'
        ) for i in range(engine_count)
    ])

    # Run the async function to embed the text in parallel
    results = asyncio.run(run_parallel_embeddings(dataset['text']))

    # Optionally, handle or display the embeddings here
    # For example:
    print(f"Results: {results}")

    output_file = os.path.join(args.output_dir, f"all_embeddings_{model_name}.npy")
    np.save(output_file, results)

[wirthual]

You can try the tqdm wrapper for gather:

See:
https://tqdm.github.io/docs/asyncio/#gather

[rangehow]

You can try the tqdm wrapper for gather:

See:
https://tqdm.github.io/docs/asyncio/#gather

thanks but it's not a good choice
(reason I have mentioned)
In a multiprocessing (or similar concurrent) scenario, to reduce the overhead of creating new tasks, an optional approach is to divide the tasks according to the number of workers (here corresponding to the engine). In this case, since the total number of worker is very small, using tqdm outside would result in a coarse-grained observation (for example, the progress bar might increment by one every few hours, and the total length of the progress bar could be very short, depending on the number of workers).

[michaelfeil]

Can you use the

‘async with engine’ call outside of the for loop - engine start and stop does take time & kills all the performance benefits.
Also your function runs them sequentially, not parallel. Use async tasks or asyncio.gather for that, no for loops.

[rangehow]

Thanks for your advice, but we may need more guide to finish multi GPU usage via infinity. Can you add some example code for this?

Also I am curious about how to do ‘async with engine’ call outside of the for loop