Spaces:
Runtime error
Runtime error
File size: 1,899 Bytes
8cedc06 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from transformers import AutoModel
class ImageEncoder(nn.Module):
"""
Image encoder using a pretrained ResNet18 backbone.
Outputs normalized embeddings.
"""
def __init__(self, embed_dim=256):
super().__init__()
self.backbone = timm.create_model(
"resnet18",
pretrained=True,
num_classes=0
)
self.proj = nn.Linear(512, embed_dim)
def forward(self, x):
x = self.backbone(x)
x = self.proj(x)
return F.normalize(x, dim=-1)
class AudioEncoder(nn.Module):
"""
CNN-based encoder for mel spectrograms.
"""
def __init__(self, embed_dim=256):
super().__init__()
self.cnn = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.AdaptiveAvgPool2d((1, 1))
)
self.proj = nn.Linear(64, embed_dim)
def forward(self, x):
x = self.cnn(x)
x = x.flatten(1)
x = self.proj(x)
return F.normalize(x, dim=-1)
class TextEncoder(nn.Module):
"""
DistilBERT-based text encoder.
"""
def __init__(self, embed_dim=256):
super().__init__()
self.backbone = AutoModel.from_pretrained(
"distilbert-base-uncased"
)
self.proj = nn.Linear(768, embed_dim)
def forward(self, input_ids, attention_mask):
output = self.backbone(
input_ids=input_ids,
attention_mask=attention_mask
)
cls = output.last_hidden_state[:, 0, :]
x = self.proj(cls)
return F.normalize(x, dim=-1) |