import torch
from transformers import AutoModel, AutoImageProcessor
from datasets import load_dataset
from safetensors.torch import save_file


ds = load_dataset("wbensvage/clothes_desc")["train"]
ds = ds.select_columns(["image"])

model_name = "google/siglip2-large-patch16-512"

model = AutoModel.from_pretrained(model_name, device_map="auto").eval()
processor = AutoImageProcessor.from_pretrained(model_name)

def encode_images(examples):
    images = examples["image"]
    inputs = processor(images=images, return_tensors="pt").to(model.device)
    with torch.no_grad():
        image_embeddings = model.get_image_features(**inputs)
    return {"vector": image_embeddings.detach().cpu()}


print(model.device)


ds = ds.map(encode_images, batched=True, batch_size=32)
ds.set_format(type="torch", columns=["vector"])

print(ds["vector"].shape)

save_file({"vectors": ds["vector"]}, "clothes_desc.safetensors")