generate_embeddings.py

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset, DistributedSampler
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
import numpy as np
from tqdm import tqdm

# Custom dataset to handle text loading
class TextDataset(Dataset):
    def __init__(self, texts):
        self.texts = texts

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

    def __getitem__(self, idx):
        return self.texts[idx]

# Collate function for data loading
def collate_fn(batch):
    return batch

# Function to save embeddings
def save_embeddings(embeddings, rank, output_dir):
    output_file = os.path.join(output_dir, f"embeddings_rank_{rank}.npy")
    os.makedirs(output_dir, exist_ok=True)
    np.save(output_file, embeddings)
    print(f"Saved embeddings for rank {rank} at {output_file}")

# Main training function
def train(rank, world_size, dataset, output_dir, batch_size=1024):
    # Initialize the process group
    dist.init_process_group(backend="nccl", init_method="env://")
    torch.cuda.set_device(rank)
    device = torch.device(f"cuda:{rank}")

    # Load the model
    model = SentenceTransformer('dunzhang/stella_en_400M_v5', trust_remote_code=True).cuda()
    model = DDP(model, device_ids=[rank])

    # Prepare data for distributed training
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=False)
    dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, collate_fn=collate_fn, pin_memory=True, num_workers=32)

    # Generate embeddings
    embeddings = []
    progress_bar = tqdm(dataloader, desc=f"Rank {rank} Progress", position=rank)
    for batch in progress_bar:
        with torch.no_grad():
            batch_embeddings = model.module.encode(batch, convert_to_tensor=True, device=device)
            embeddings.append(batch_embeddings.cpu().numpy())
    embeddings = np.vstack(embeddings)

    # Save embeddings for this rank
    save_embeddings(embeddings, rank, output_dir)

    # Finalize
    dist.barrier()
    dist.destroy_process_group()

# Entry point for torchrun
def main(local_rank, world_size, dataset_name="ai-practicum-group/post2geo-dataset", output_dir="embeddings_output"):
    # Load dataset
    dataset = load_dataset(dataset_name)["train"]
    texts = dataset["text"]
    text_dataset = TextDataset(texts)

    # Launch distributed training
    train(local_rank, world_size, text_dataset, output_dir)

    # Combine embeddings if rank 0
    if local_rank == 0:
        all_embeddings = []
        for rank in range(world_size):
            embeddings_file = os.path.join(output_dir, f"embeddings_rank_{rank}.npy")
            all_embeddings.append(np.load(embeddings_file))
        all_embeddings = np.vstack(all_embeddings)

        # Save the combined dataset
        dataset = dataset.add_column("embeddings", all_embeddings.tolist())
        dataset.save_to_disk("embedded_dataset")
        print("Combined embeddings saved to embedded_dataset")

if __name__ == "__main__":
    # Entry point for torchrun
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    main(local_rank, world_size)