"""
This example starts multiple processes (1 per GPU), which encode
sentences in parallel. This gives a near linear speed-up
when encoding large text collections.
"""
from tqdm import tqdm
import logging
import json
import torch
from sentence_transformers import LoggingHandler, SentenceTransformer

logging.basicConfig(
    format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO, handlers=[LoggingHandler()]
)

# Important, you need to shield your code with if __name__. Otherwise, CUDA runs into issues when spawning new processes.
if __name__ == "__main__":
    # Create a large list of 100k sentences
    f = open("/home/aiscuser/fhw/data/qwq_python_selected.json","r+")
    lines = f.readlines()
    sentences = []
    for line in tqdm(lines):
        d= json.loads(line)
        sentences.append(d["instruction"])


    # Define the model
    model = SentenceTransformer("/home/aiscuser/fhw/model_weights/all-roberta-large-v1")

    # Start the multi-process pool on all available CUDA devices
    pool = model.start_multi_process_pool(["cuda:0", "cuda:0", "cuda:0", "cuda:0", "cuda:0", "cuda:0", "cuda:0", "cuda:0", "cuda:0", "cuda:0", "cuda:1", "cuda:1", "cuda:1", "cuda:1", "cuda:1", "cuda:1", "cuda:1", "cuda:1", "cuda:1", "cuda:1", "cuda:2", "cuda:2", "cuda:2", "cuda:2", "cuda:2", "cuda:2", "cuda:2", "cuda:2", "cuda:2", "cuda:2", "cuda:3", "cuda:3", "cuda:3", "cuda:3", "cuda:3", "cuda:3", "cuda:3", "cuda:3", "cuda:3", "cuda:3", "cuda:4", "cuda:4", "cuda:4", "cuda:4", "cuda:4", "cuda:4", "cuda:4", "cuda:4", "cuda:4", "cuda:4", "cuda:5", "cuda:5", "cuda:5", "cuda:5", "cuda:5", "cuda:5", "cuda:5", "cuda:5", "cuda:5", "cuda:5", "cuda:6", "cuda:6", "cuda:6", "cuda:6", "cuda:6", "cuda:6", "cuda:6", "cuda:6", "cuda:6", "cuda:6", "cuda:7", "cuda:7", "cuda:7", "cuda:7", "cuda:7", "cuda:7", "cuda:7", "cuda:7", "cuda:7", "cuda:7", ])

    # Compute the embeddings using the multi-process pool
    emb = model.encode_multi_process(sentences, pool)
    
    print("Embeddings computed. Shape:", emb.shape)

    # Optional: Stop the processes in the pool
    model.stop_multi_process_pool(pool)
    torch.save(emb, "/home/aiscuser/fhw/embeddings/qwq_ins_embeddings.pt", pickle_protocol=4)