"""High-level Forge2Vec API and vector arithmetic objects.""" from __future__ import annotations from dataclasses import dataclass, field, replace from pathlib import Path from typing import Any, Iterable, Iterator, Mapping, Optional, Sequence, Union import numpy as np import torch from PIL import Image from safetensors.torch import load_file from transformers import AutoTokenizer from .features import genre_indices, metadata_vector from .modeling import UnifiedAttentionForge2Vec Poster = Union[Image.Image, np.ndarray, torch.Tensor] @dataclass(frozen=True) class ForgeItem: """An anime or manga profile accepted by Forge2Vec.""" title: str = "" native_title: str = "" synonyms: Sequence[str] = field(default_factory=tuple) synopsis: str = "" synopsis_ua: str = "" genres: Sequence[str] = field(default_factory=tuple) content_type: str = "anime" year: Optional[int] = None score: Optional[float] = None poster: Optional[Poster] = field(default=None, repr=False, compare=False) id: Optional[Union[str, int]] = None @classmethod def from_value(cls, value: Union["ForgeItem", Mapping[str, Any]]) -> "ForgeItem": if isinstance(value, cls): return value if not isinstance(value, Mapping): raise TypeError("item must be a ForgeItem or a mapping") aliases = { "ua_title": "title", "en_title": "title", "original_title": "native_title", "alternate_names": "synonyms", "ua_description": "synopsis_ua", "en_description": "synopsis", "type": "content_type", } normalized = dict(value) for old, new in aliases.items(): if new not in normalized and old in normalized: normalized[new] = normalized[old] fields = cls.__dataclass_fields__ return cls(**{key: value for key, value in normalized.items() if key in fields}) class ForgeVector: """A Forge2Vec embedding supporting ordinary vector arithmetic.""" __array_priority__ = 1000 def __init__( self, values: Union[np.ndarray, torch.Tensor, Sequence[float]], *, catalogue: Optional["ForgeCatalogue"] = None, excluded_indices: Iterable[int] = (), ) -> None: array = np.asarray(values, dtype=np.float32).reshape(-1) if array.shape != (256,): raise ValueError(f"ForgeVector must have shape (256,), received {array.shape}") self._values = array self._catalogue = catalogue self._excluded_indices = frozenset(excluded_indices) @property def values(self) -> np.ndarray: return self._values.copy() @property def shape(self) -> tuple[int, ...]: return self._values.shape def numpy(self) -> np.ndarray: return self.values def tensor(self, device: Optional[Union[str, torch.device]] = None) -> torch.Tensor: return torch.from_numpy(self._values.copy()).to(device=device) def normalized(self) -> "ForgeVector": norm = float(np.linalg.norm(self._values)) if norm <= 1e-9: raise ValueError("cannot normalize a zero vector") return self._new(self._values / norm) def find(self, limit: int = 10) -> "ForgeResults": if self._catalogue is None: raise ValueError("this vector is not attached to a catalogue") return self._catalogue.find(self, limit=limit) def _new(self, values, other: Optional["ForgeVector"] = None) -> "ForgeVector": catalogue = self._catalogue excluded = self._excluded_indices if other is not None: if catalogue is None: catalogue = other._catalogue elif other._catalogue is not None and other._catalogue is not catalogue: catalogue = None excluded = excluded | other._excluded_indices return ForgeVector(values, catalogue=catalogue, excluded_indices=excluded) def __add__(self, other: "ForgeVector") -> "ForgeVector": if not isinstance(other, ForgeVector): return NotImplemented return self._new(self._values + other._values, other) def __sub__(self, other: "ForgeVector") -> "ForgeVector": if not isinstance(other, ForgeVector): return NotImplemented return self._new(self._values - other._values, other) def __mul__(self, scalar: float) -> "ForgeVector": return self._new(self._values * float(scalar)) def __rmul__(self, scalar: float) -> "ForgeVector": return self * scalar def __truediv__(self, scalar: float) -> "ForgeVector": if float(scalar) == 0.0: raise ZeroDivisionError("cannot divide a ForgeVector by zero") return self._new(self._values / float(scalar)) def __neg__(self) -> "ForgeVector": return self._new(-self._values) def __array__(self, dtype=None) -> np.ndarray: return np.asarray(self._values, dtype=dtype) def __repr__(self) -> str: return f"ForgeVector(shape={self.shape}, norm={np.linalg.norm(self._values):.4f})" @dataclass(frozen=True) class ForgeMatch: rank: int item: ForgeItem vector: ForgeVector similarity: float class ForgeResults(Sequence[ForgeMatch]): def __init__(self, matches: Sequence[ForgeMatch]) -> None: self._matches = tuple(matches) def __getitem__(self, index): return self._matches[index] def __len__(self) -> int: return len(self._matches) def __iter__(self) -> Iterator[ForgeMatch]: return iter(self._matches) def __repr__(self) -> str: lines = ["ForgeResults("] lines.extend( f" {match.rank}. {match.item.title} ({match.similarity:.4f})" for match in self._matches ) return "\n".join((*lines, ")")) class ForgeCatalogue: """An in-memory cosine index for anime and manga vectors.""" def __init__(self, model: "Forge2Vec", items: Sequence[ForgeItem], embeddings: np.ndarray) -> None: self.model = model self.items = tuple(items) values = np.asarray(embeddings, dtype=np.float32) norms = np.linalg.norm(values, axis=1, keepdims=True) self.embeddings = values / np.maximum(norms, 1e-9) self._titles: dict[str, int] = {} for index, item in enumerate(self.items): for title in (item.title, item.native_title, *item.synonyms): if title: self._titles.setdefault(title.casefold().strip(), index) def vec(self, title_or_index: Union[str, int]) -> ForgeVector: if isinstance(title_or_index, str): key = title_or_index.casefold().strip() if key not in self._titles: raise KeyError(f"title is not present in the catalogue: {title_or_index}") index = self._titles[key] else: index = int(title_or_index) return ForgeVector(self.embeddings[index], catalogue=self, excluded_indices=(index,)) def find(self, query: Union[ForgeVector, ForgeItem, Mapping[str, Any]], limit: int = 10) -> ForgeResults: vector = query if isinstance(query, ForgeVector) else self.model.vec(query) normalized = vector.normalized()._values scores = self.embeddings @ normalized order = np.argsort(scores)[::-1] selected = [index for index in order if index not in vector._excluded_indices][:limit] matches = [ ForgeMatch( rank=rank, item=self.items[index], vector=ForgeVector(self.embeddings[index], catalogue=self, excluded_indices=(index,)), similarity=float(scores[index]), ) for rank, index in enumerate(selected, 1) ] return ForgeResults(matches) class Forge2Vec: """Load and run hikka-forge2vec.""" default_model_id = "Lorg0n/hikka-forge2vec" def __init__( self, model_id_or_path: Union[str, Path] = default_model_id, *, device: Optional[Union[str, torch.device]] = None, revision: Optional[str] = None, cache_dir: Optional[Union[str, Path]] = None, ) -> None: self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) root = Path(model_id_or_path) if not root.is_dir(): from huggingface_hub import snapshot_download root = Path(snapshot_download( repo_id=str(model_id_or_path), revision=revision, cache_dir=cache_dir, allow_patterns=("config.json", "model.safetensors", "assets/tokenizer/*"), )) import json self.config = json.loads((root / "config.json").read_text(encoding="utf-8")) tokenizer_path = root / "assets" / "tokenizer" self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, local_files_only=True) self.model = UnifiedAttentionForge2Vec( str(tokenizer_path), max_style_weight=float(self.config["max_style_weight"]) ) incompatible = self.model.load_state_dict( load_file(str(root / "model.safetensors"), device="cpu"), strict=False ) if incompatible.missing_keys or incompatible.unexpected_keys: raise RuntimeError( f"incompatible model artifact: missing={incompatible.missing_keys}, " f"unexpected={incompatible.unexpected_keys}" ) self.model.to(self.device).eval() @staticmethod def _poster_tensor(poster: Poster) -> torch.Tensor: if isinstance(poster, Image.Image): image = poster.convert("RGB").resize((224, 224), Image.Resampling.BICUBIC) tensor = torch.from_numpy(np.asarray(image, dtype=np.float32).copy()).permute(2, 0, 1) else: tensor = torch.as_tensor(poster).detach().to(dtype=torch.float32, device="cpu") if tensor.ndim != 3: raise ValueError("poster tensor or array must have three dimensions") if tensor.shape[0] not in (1, 3, 4) and tensor.shape[-1] in (1, 3, 4): tensor = tensor.permute(2, 0, 1) if tensor.shape[0] == 1: tensor = tensor.expand(3, -1, -1) elif tensor.shape[0] == 4: tensor = tensor[:3] tensor = torch.nn.functional.interpolate( tensor.unsqueeze(0), size=(224, 224), mode="bicubic", align_corners=False ).squeeze(0) if tensor.max() > 1.0: tensor = tensor / 255.0 if tensor.min() >= 0.0: tensor = tensor * 2.0 - 1.0 return tensor.clamp(-1.0, 1.0) def _tokenize(self, texts: Sequence[str]) -> tuple[torch.Tensor, torch.Tensor]: tokens = self.tokenizer( list(texts), padding=True, truncation=True, max_length=int(self.config["text_max_length"]), return_tensors="pt", ) return tokens["input_ids"].to(self.device), tokens["attention_mask"].to(self.device) @torch.inference_mode() def vecs( self, items: Iterable[Union[ForgeItem, Mapping[str, Any]]], *, batch_size: int = 32, ) -> list[ForgeVector]: profiles = [ForgeItem.from_value(item) for item in items] output: list[ForgeVector] = [] for start in range(0, len(profiles), batch_size): batch = profiles[start:start + batch_size] descriptions_ua = [item.synopsis_ua or item.synopsis for item in batch] descriptions_en = [item.synopsis or item.synopsis_ua for item in batch] titles = [", ".join(filter(None, (item.title, item.native_title, *item.synonyms))) for item in batch] ua_ids, ua_mask = self._tokenize(descriptions_ua) en_ids, en_mask = self._tokenize(descriptions_en) title_ids, title_mask = self._tokenize(titles) ua = self.model.encode_text(ua_ids, ua_mask) en = self.model.encode_text(en_ids, en_mask) title_vectors = self.model.encode_text(title_ids, title_mask) genres = torch.tensor([genre_indices(list(item.genres)) for item in batch], device=self.device) metadata = torch.tensor([ metadata_vector(item.content_type, item.year, item.score) for item in batch ], dtype=torch.float32, device=self.device) mask = torch.tensor([ [float(bool(item.synopsis or item.synopsis_ua)), float(bool(item.title or item.native_title or item.synonyms)), float(bool(item.genres)), 1.0, float(item.poster is not None)] for item in batch ], dtype=torch.float32, device=self.device) posters = torch.zeros((len(batch), 3, 224, 224), device=self.device) for index, item in enumerate(batch): if item.poster is not None: posters[index] = self._poster_tensor(item.poster).to(self.device) embeddings = self.model(ua, en, title_vectors, genres, metadata, mask, posters) output.extend(ForgeVector(row) for row in embeddings.cpu().numpy()) return output def vec(self, item: Optional[Union[ForgeItem, Mapping[str, Any]]] = None, **fields: Any) -> ForgeVector: if item is not None and fields: profile = replace(ForgeItem.from_value(item), **fields) elif item is not None: profile = ForgeItem.from_value(item) else: profile = ForgeItem(**fields) return self.vecs([profile])[0] def catalogue( self, items: Iterable[Union[ForgeItem, Mapping[str, Any]]], *, batch_size: int = 32, ) -> ForgeCatalogue: profiles = [ForgeItem.from_value(item) for item in items] vectors = self.vecs(profiles, batch_size=batch_size) embeddings = np.stack([vector._values for vector in vectors]) return ForgeCatalogue(self, profiles, embeddings)