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
82
83
84
85
86
87
88
89
90
|
from __future__ import annotations
import torch
from torch import nn
import torch.nn.functional as F
class Bridge(nn.Module):
def __init__(
self,
vision_dim: int,
text_dim: int,
hidden_dim: int = 1536,
linear: bool = False,
):
super().__init__()
self.vision_dim = vision_dim
self.text_dim = text_dim
self.hidden_dim = hidden_dim
self.linear = linear
if linear:
self.network = nn.Linear(vision_dim, text_dim, bias=False)
else:
self.network = nn.Sequential(
nn.LayerNorm(vision_dim),
nn.Linear(vision_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, text_dim),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.normalize(self.network(x).float(), dim=-1)
def config(self) -> dict:
return {
"vision_dim": self.vision_dim,
"text_dim": self.text_dim,
"hidden_dim": self.hidden_dim,
"linear": self.linear,
}
class PrefixAdapter(nn.Module):
def __init__(
self,
semantic_dim: int,
lm_dim: int,
prefix_length: int = 8,
hidden_dim: int = 2048,
):
super().__init__()
self.semantic_dim = semantic_dim
self.lm_dim = lm_dim
self.prefix_length = prefix_length
self.hidden_dim = hidden_dim
self.network = nn.Sequential(
nn.LayerNorm(semantic_dim),
nn.Linear(semantic_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, prefix_length * lm_dim),
)
def forward(self, semantic: torch.Tensor) -> torch.Tensor:
prefix = self.network(semantic.float())
return prefix.view(-1, self.prefix_length, self.lm_dim)
def config(self) -> dict:
return {
"semantic_dim": self.semantic_dim,
"lm_dim": self.lm_dim,
"prefix_length": self.prefix_length,
"hidden_dim": self.hidden_dim,
}
def load_bridge(path: str, device: str = "cpu") -> tuple[Bridge, dict]:
state = torch.load(path, map_location="cpu", weights_only=False)
model = Bridge(**state["config"])
model.load_state_dict(state["state_dict"])
model.to(device).eval()
return model, state
def load_prefix(path: str, device: str = "cpu") -> tuple[PrefixAdapter, dict]:
state = torch.load(path, map_location="cpu", weights_only=False)
model = PrefixAdapter(**state["config"])
model.load_state_dict(state["state_dict"])
model.to(device).eval()
return model, state
|
