File size: 1,840 Bytes
50611ef | 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 | """Keyword-based text encoder for conditioning the diffusion model."""
import csv
from pathlib import Path
import torch
import torch.nn as nn
def build_vocab(metadata_csv: Path, min_count: int = 5) -> list[str]:
"""Build vocabulary from metadata CSV keywords, filtering rare tokens."""
counts: dict[str, int] = {}
with open(metadata_csv) as f:
reader = csv.DictReader(f)
for row in reader:
for kw in row["keywords"].split(","):
kw = kw.strip().lower()
if kw:
counts[kw] = counts.get(kw, 0) + 1
vocab = sorted(kw for kw, c in counts.items() if c >= min_count)
return vocab
class KeywordEncoder(nn.Module):
"""Encodes a set of keyword indices into a fixed-size conditioning vector."""
def __init__(
self,
vocab_size: int,
embed_dim: int = 64,
cond_dim: int = 256,
) -> None:
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.proj = nn.Linear(embed_dim, cond_dim)
self.null_embedding = nn.Parameter(torch.randn(cond_dim))
def forward(self, token_ids: list[list[int]], device: torch.device) -> torch.Tensor:
"""Encode batch of keyword lists into conditioning vectors.
Args:
token_ids: List of lists of vocab indices per sample.
device: Target device.
Returns:
(batch, cond_dim) conditioning tensor.
"""
batch_embs = []
for ids in token_ids:
if len(ids) == 0:
batch_embs.append(self.null_embedding)
else:
idx = torch.tensor(ids, device=device)
emb = self.embedding(idx).mean(dim=0)
batch_embs.append(self.proj(emb))
return torch.stack(batch_embs)
|