Reference for ultralytics/models/yolo/world/train.py

def build_dataset(self, img_path: str, mode: str = "train", batch: Optional[int] = None):
    """
    Build YOLO Dataset for training or validation.

    Args:
        img_path (str): Path to the folder containing images.
        mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
        batch (int, optional): Size of batches, this is for `rect`.

    Returns:
        (Any): YOLO dataset configured for training or validation.
    """
    gs = max(int(de_parallel(self.model).stride.max() if self.model else 0), 32)
    dataset = build_yolo_dataset(
        self.args, img_path, batch, self.data, mode=mode, rect=mode == "val", stride=gs, multi_modal=mode == "train"
    )
    if mode == "train":
        self.set_text_embeddings([dataset], batch)  # cache text embeddings to accelerate training
    return dataset

def generate_text_embeddings(self, texts: List[str], batch: int, cache_dir: Path) -> Dict[str, torch.Tensor]:
    """
    Generate text embeddings for a list of text samples.

    Args:
        texts (List[str]): List of text samples to encode.
        batch (int): Batch size for processing.
        cache_dir (Path): Directory to save/load cached embeddings.

    Returns:
        (Dict[str, torch.Tensor]): Dictionary mapping text samples to their embeddings.
    """
    model = "clip:ViT-B/32"
    cache_path = cache_dir / f"text_embeddings_{model.replace(':', '_').replace('/', '_')}.pt"
    if cache_path.exists():
        LOGGER.info(f"Reading existed cache from '{cache_path}'")
        txt_map = torch.load(cache_path)
        if sorted(txt_map.keys()) == sorted(texts):
            return txt_map
    LOGGER.info(f"Caching text embeddings to '{cache_path}'")
    assert self.model is not None
    txt_feats = self.model.get_text_pe(texts, batch, cache_clip_model=False)
    txt_map = dict(zip(texts, txt_feats.squeeze(0)))
    torch.save(txt_map, cache_path)
    return txt_map

def get_model(self, cfg=None, weights: Optional[str] = None, verbose: bool = True) -> WorldModel:
    """
    Return WorldModel initialized with specified config and weights.

    Args:
        cfg (Dict[str, Any] | str, optional): Model configuration.
        weights (str, optional): Path to pretrained weights.
        verbose (bool): Whether to display model info.

    Returns:
        (WorldModel): Initialized WorldModel.
    """
    # NOTE: This `nc` here is the max number of different text samples in one image, rather than the actual `nc`.
    # NOTE: Following the official config, nc hard-coded to 80 for now.
    model = WorldModel(
        cfg["yaml_file"] if isinstance(cfg, dict) else cfg,
        ch=self.data["channels"],
        nc=min(self.data["nc"], 80),
        verbose=verbose and RANK == -1,
    )
    if weights:
        model.load(weights)
    self.add_callback("on_pretrain_routine_end", on_pretrain_routine_end)

    return model

def preprocess_batch(self, batch: Dict[str, Any]) -> Dict[str, Any]:
    """Preprocess a batch of images and text for YOLOWorld training."""
    batch = DetectionTrainer.preprocess_batch(self, batch)

    # Add text features
    texts = list(itertools.chain(*batch["texts"]))
    txt_feats = torch.stack([self.text_embeddings[text] for text in texts]).to(self.device)
    txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True)
    batch["txt_feats"] = txt_feats.reshape(len(batch["texts"]), -1, txt_feats.shape[-1])
    return batch

def set_text_embeddings(self, datasets: List[Any], batch: Optional[int]) -> None:
    """
    Set text embeddings for datasets to accelerate training by caching category names.

    This method collects unique category names from all datasets, then generates and caches text embeddings
    for these categories to improve training efficiency.

    Args:
        datasets (List[Any]): List of datasets from which to extract category names.
        batch (int | None): Batch size used for processing.

    Notes:
        This method collects category names from datasets that have the 'category_names' attribute,
        then uses the first dataset's image path to determine where to cache the generated text embeddings.
    """
    text_embeddings = {}
    for dataset in datasets:
        if not hasattr(dataset, "category_names"):
            continue
        text_embeddings.update(
            self.generate_text_embeddings(
                list(dataset.category_names), batch, cache_dir=Path(dataset.img_path).parent
            )
        )
    self.text_embeddings = text_embeddings