48 lines
1.7 KiB
Python
48 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from data.schema import ENTITIES, PREDICATES
|
|
from models.frozen_backbone import build_backbone
|
|
from models.slot_extractor import SlotExtractor
|
|
|
|
|
|
class EvidenceExtractorModel(nn.Module):
|
|
def __init__(self, config: dict[str, Any], backbone_name: str | None = None):
|
|
super().__init__()
|
|
model_config = config["model"]
|
|
selected = backbone_name or model_config["backbone"]
|
|
self.backbone_name = selected
|
|
self.backbone = build_backbone(selected, model_config["clip_model"])
|
|
common = {
|
|
"state_dim": model_config["state_dim"],
|
|
"num_layers": model_config["extractor_layers"],
|
|
"num_heads": model_config["num_heads"],
|
|
"dropout": model_config["dropout"],
|
|
"num_subjects": len(ENTITIES),
|
|
"num_predicates": len(PREDICATES),
|
|
"num_objects": len(ENTITIES),
|
|
}
|
|
self.visual_extractor = SlotExtractor(
|
|
input_dim=self.backbone.vision_dim,
|
|
num_slots=model_config["num_visual_slots"],
|
|
**common,
|
|
)
|
|
self.text_extractor = SlotExtractor(
|
|
input_dim=self.backbone.text_dim,
|
|
num_slots=model_config["num_text_slots"],
|
|
**common,
|
|
)
|
|
|
|
def forward(self, images: torch.Tensor, rules_text: list[str]) -> dict[str, Any]:
|
|
visual_tokens, visual_mask = self.backbone.encode_image(images)
|
|
text_tokens, text_mask = self.backbone.encode_text(rules_text)
|
|
return {
|
|
"visual": self.visual_extractor(visual_tokens, visual_mask),
|
|
"text": self.text_extractor(text_tokens, text_mask),
|
|
}
|
|
|