Используй API datasets для использования CoCo датасета, чтобы я не загружал с сайта. import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

from transformers import BertTokenizer, BertModel
from torchvision.models import resnet50
import torch.optim as optim


import torch.nn.functional as F



# Используем стандартные аугментации для изображений
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

# Загружаем набор данных COCO или LAION (пример COCO)
dataset = datasets.CocoCaptions(
    root='path/to/images',
    annFile='path/to/annotations/captions.json',
    transform=transform
)

dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# Инициализация текстового энкодера (предобученный BERT)
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
text_encoder = BertModel.from_pretrained('bert-base-uncased')

# Инициализация визуального энкодера (ResNet-50)
image_encoder = resnet50(pretrained=True)
image_encoder.fc = torch.nn.Identity()  # Убираем финальный классификатор



def contrastive_loss(image_embeddings, text_embeddings, temperature=0.07):
    # Косинусное сходство
    image_embeddings = F.normalize(image_embeddings, dim=-1)
    text_embeddings = F.normalize(text_embeddings, dim=-1)

    logits = torch.matmul(image_embeddings, text_embeddings.T) / temperature
    labels = torch.arange(len(logits)).to(logits.device)

    # Кросс-энтропийная функция потерь
    loss_img_to_txt = F.cross_entropy(logits, labels)
    loss_txt_to_img = F.cross_entropy(logits.T, labels)

    return (loss_img_to_txt + loss_txt_to_img) / 2



# Параметры модели
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
text_encoder = text_encoder.to(device)
image_encoder = image_encoder.to(device)

# Оптимизатор
optimizer = optim.Adam(list(text_encoder.parameters()) + list(image_encoder.parameters()), lr=1e-4)

# Цикл обучения
for epoch in range(10):  # Тренируем в течение 10 эпох
    for images, captions in dataloader:
        # Преобразование текста
        inputs = tokenizer(captions, padding=True, truncation=True, return_tensors="pt")
        input_ids = inputs.input_ids.to(device)
        attention_mask = inputs.attention_mask.to(device)

        # Генерация эмбеддингов для изображений и текста
        image_embeddings = image_encoder(images.to(device))
        text_embeddings = text_encoder(input_ids, attention_mask=attention_mask).pooler_output

        # Контрастивная функция потерь
        loss = contrastive_loss(image_embeddings, text_embeddings)

        # Оптимизация
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch}, Loss: {loss.item()}")


def retrieve_image(query, dataset, tokenizer, text_encoder, image_encoder, top_k=5):
    text_input = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        text_embedding = text_encoder(text_input.input_ids.to(device),
                                      attention_mask=text_input.attention_mask.to(device)).pooler_output

    image_embeddings = []
    for img, _ in dataset:
        with torch.no_grad():
            img_emb = image_encoder(img.unsqueeze(0).to(device))
            image_embeddings.append(img_emb)

    image_embeddings = torch.cat(image_embeddings)
    scores = torch.matmul(text_embedding, image_embeddings.T)

    # Выводим топ-k изображений
    top_k_indices = torch.topk(scores, k=top_k).indices.squeeze().tolist()
    return [dataset[i][0] for i in top_k_indices]


# Пример использования
query = "A cat sitting on a couch"
top_images = retrieve_image(query, dataset, tokenizer, text_encoder, image_encoder)