@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1f5d13ab9f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 14:45 UTC

[New Model]: jinaai/jina-embeddings-v3 (#16120)

wang.yuqi committed Apr 8, 2025 at 23:39 UTC 1f5d13ab9f1232751a09cdb6b1cdbb4687393833

6 files changed +297 -86

examples/offline_inference/embed_jina_embeddings_v3.py new

+50

@@ -0,0 +1,50 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +from argparse import Namespace

4 +

5 +from vllm import LLM, EngineArgs

6 +from vllm.utils import FlexibleArgumentParser

7 +

8 +

9 +def main(args: Namespace):

10 + # Sample prompts.

11 + prompts = [

12 + "Follow the white rabbit.", # English

13 + "Sigue al conejo blanco.", # Spanish

14 + "Suis le lapin blanc.", # French

15 + "跟着白兔走。", # Chinese

16 + "اتبع الأرنب الأبيض.", # Arabic

17 + "Folge dem weißen Kaninchen.", # German

18 + ]

19 +

20 + # Create an LLM.

21 + # You should pass task="embed" for embedding models

22 + model = LLM(**vars(args))

23 +

24 + # Generate embedding. The output is a list of EmbeddingRequestOutputs.

25 + # Only text matching task is supported for now. See #16120

26 + outputs = model.embed(prompts)

27 +

28 + # Print the outputs.

29 + print("\nGenerated Outputs:")

30 + print("Only text matching task is supported for now. See #16120")

31 + print("-" * 60)

32 + for prompt, output in zip(prompts, outputs):

33 + embeds = output.outputs.embedding

34 + embeds_trimmed = ((str(embeds[:16])[:-1] +

35 + ", ...]") if len(embeds) > 16 else embeds)

36 + print(f"Prompt: {prompt!r} \n"

37 + f"Embeddings for text matching: {embeds_trimmed} "

38 + f"(size={len(embeds)})")

39 + print("-" * 60)

40 +

41 +

42 +if __name__ == "__main__":

43 + parser = FlexibleArgumentParser()

44 + parser = EngineArgs.add_cli_args(parser)

45 + # Set example specific arguments

46 + parser.set_defaults(model="jinaai/jina-embeddings-v3",

47 + task="embed",

48 + trust_remote_code=True)

49 + args = parser.parse_args()

50 + main(args)

tests/conftest.py

+3 -2

@@ -671,8 +671,9 @@ class HfRunner:

671 return [(output_ids, output_str, output_logprobs)

672 for output_ids, output_str, output_logprobs in outputs]

673

674 - def encode(self, prompts: list[str]) -> list[list[torch.Tensor]]:

675 - return self.model.encode(prompts)

674 + def encode(self, prompts: list[str], *args,

675 + **kwargs) -> list[list[torch.Tensor]]:

676 + return self.model.encode(prompts, *args, **kwargs)

677

678 def predict(self, prompts: list[list[str]]) -> torch.Tensor:

679 return self.model.predict(prompts, convert_to_tensor=True)

tests/models/embedding/language/test_jina.py renamed

+61 -3

@@ -2,13 +2,15 @@

2 # ruff: noqa: E501

3 """Compare the scoring outputs of HF and vLLM models.

5 -Run `pytest tests/models/embedding/language/test_jina_reranker_v2.py`.

5 +Run `pytest tests/models/embedding/language/test_jina.py`.

6 """

7 import math

9 import pytest

11 -MODELS = [

11 +from tests.models.embedding.utils import check_embeddings_close

12 +

13 +SCORING_MODELS = [

14 "jinaai/jina-reranker-v2-base-multilingual", # Roberta

15 ]

@@ -27,8 +29,21 @@ TEXTS_2 = [

29 "新しいメイクのトレンドは鮮やかな色と革新的な技術に焦点を当てています",

30 ]

32 +EMBEDDING_MODELS = [

33 + "jinaai/jina-embeddings-v3",

34 +]

35 +

36 +EMBEDDING_PROMPTS = [

37 + "Follow the white rabbit.", # English

38 + "Sigue al conejo blanco.", # Spanish

39 + "Suis le lapin blanc.", # French

40 + "跟着白兔走。", # Chinese

41 + "اتبع الأرنب الأبيض.", # Arabic

42 + "Folge dem weißen Kaninchen.", # German

43 +]

44 +

31 -@pytest.fixture(scope="module", params=MODELS)

46 +@pytest.fixture(scope="module", params=SCORING_MODELS)

47 def model_name(request):

48 yield request.param

@@ -68,3 +83,46 @@ def test_llm_1_to_N(vllm_runner, hf_runner, model_name, dtype: str):

84 assert math.isclose(hf_outputs[0], vllm_outputs[0], rel_tol=0.01)

85 assert math.isclose(hf_outputs[1], vllm_outputs[1], rel_tol=0.01)

86 +

87 +

88 +@pytest.fixture(scope="module", params=EMBEDDING_MODELS)

89 +def emb_model_name(request):

90 + yield request.param

91 +

92 +

93 +def test_is_matryoshka(vllm_runner, emb_model_name):

94 + with vllm_runner(emb_model_name, task="embed",

95 + max_model_len=None) as vllm_model:

96 + assert vllm_model.model.llm_engine.model_config.is_matryoshka

97 +

98 +

99 +@pytest.mark.parametrize("model", EMBEDDING_MODELS)

100 +@pytest.mark.parametrize("dtype", ["half"])

101 +def test_embeddings(

102 + hf_runner,

103 + vllm_runner,

104 + model,

105 + dtype: str,

106 + monkeypatch,

107 +) -> None:

108 +

109 + example_prompts = EMBEDDING_PROMPTS

110 +

111 + with hf_runner(

112 + model,

113 + dtype=dtype,

114 + is_sentence_transformer=True,

115 + ) as hf_model:

116 + hf_outputs = hf_model.encode(example_prompts, task="text-matching")

117 +

118 + with vllm_runner(model, task="embed", dtype=dtype,

119 + max_model_len=None) as vllm_model:

120 + vllm_outputs = vllm_model.encode(example_prompts)

121 +

122 + check_embeddings_close(

123 + embeddings_0_lst=hf_outputs,

124 + embeddings_1_lst=vllm_outputs,

125 + name_0="hf",

126 + name_1="vllm",

127 + tol=1e-2,

128 + )

vllm/config.py

@@ -1130,6 +1130,11 @@ class ModelConfig:

1130 architectures = getattr(self.hf_config, "architectures", [])

1131 return ModelRegistry.is_v1_compatible(architectures)

1132

1133 + @property

1134 + def is_matryoshka(self) -> bool:

1135 + return (hasattr(self.hf_config, "matryoshka_dimensions")

1136 + or getattr(self.hf_config, "is_matryoshka", False))

1137 +

1138

1139 class CacheConfig:

1140 """Configuration for the KV cache.

vllm/model_executor/models/bert.py

+49 -17

@@ -18,6 +18,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,

18 from vllm.model_executor.layers.pooler import (CrossEncodingPooler, Pooler,

19 PoolingType)

20 from vllm.model_executor.layers.quantization import QuantizationConfig

21 +from vllm.model_executor.layers.rotary_embedding import get_rope

22 from vllm.model_executor.layers.vocab_parallel_embedding import (

23 VocabParallelEmbedding)

24 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

@@ -38,19 +39,24 @@ class BertEmbedding(nn.Module):

39 self.size = config.hidden_size

40 self.word_embeddings = VocabParallelEmbedding(config.vocab_size,

41 config.hidden_size)

41 - self.position_embeddings = VocabParallelEmbedding(

42 - config.max_position_embeddings, config.hidden_size)

42 +

43 self.token_type_embeddings = VocabParallelEmbedding(

44 config.type_vocab_size, config.hidden_size)

45 self.LayerNorm = nn.LayerNorm(config.hidden_size,

46 eps=config.layer_norm_eps)

47 - self.position_ids = nn.Parameter(

48 - torch.empty((1, config.max_position_embeddings)), )

48 self.position_embedding_type = config.position_embedding_type

51 - if self.position_embedding_type != "absolute":

52 - raise ValueError("Only 'absolute' position_embedding_type" +

53 - " is supported")

49 + if self.position_embedding_type == "absolute":

50 + self.position_embeddings = VocabParallelEmbedding(

51 + config.max_position_embeddings, config.hidden_size)

52 + self.position_ids = nn.Parameter(

53 + torch.empty((1, config.max_position_embeddings)), )

54 + elif self.position_embedding_type == "rotary":

55 + self.position_embeddings = None

56 + self.position_ids = None

57 + else:

58 + raise ValueError("Only 'absolute' and 'rotary' " +

59 + "position_embedding_type is supported")

61 def forward(

62 self,

@@ -64,9 +70,6 @@ class BertEmbedding(nn.Module):

70 # Input embeddings.

71 inputs_embeds = self.word_embeddings(input_ids)

67 - # Position embeddings.

68 - position_embeddings = self.position_embeddings(position_ids)

69 -

73 if token_type_ids is None:

74 token_type_ids = torch.zeros(input_shape,

75 dtype=torch.long,

@@ -74,7 +77,12 @@ class BertEmbedding(nn.Module):

78 token_type_embeddings = self.token_type_embeddings(token_type_ids)

77 - embeddings = inputs_embeds + token_type_embeddings + position_embeddings

80 + embeddings = inputs_embeds + token_type_embeddings

81 +

82 + if self.position_embedding_type == "absolute":

83 + position_embeddings = self.position_embeddings(position_ids)

84 + embeddings += position_embeddings

85 +

86 embeddings = self.LayerNorm(embeddings)

87 return embeddings

@@ -98,7 +106,10 @@ class BertPooler(nn.Module):

106 @support_torch_compile

107 class BertEncoder(nn.Module):

108

101 - def __init__(self, vllm_config: VllmConfig, prefix: str = ""):

109 + def __init__(self,

110 + vllm_config: VllmConfig,

111 + rotary_kwargs: Optional[dict] = None,

112 + prefix: str = ""):

113 super().__init__()

114 config = vllm_config.model_config.hf_config

115 cache_config = vllm_config.cache_config

@@ -107,16 +118,18 @@ class BertEncoder(nn.Module):

118 BertLayer(config=config,

119 cache_config=cache_config,

120 quant_config=quant_config,

121 + rotary_kwargs=rotary_kwargs,

122 prefix=f"{prefix}.layer.{layer_idx}")

123 for layer_idx in range(config.num_hidden_layers)

124 ])

125

126 def forward(

127 self,

128 + positions: torch.Tensor,

129 hidden_states: torch.Tensor,

130 ) -> torch.Tensor:

131 for layer in self.layer:

119 - hidden_states = layer(hidden_states)

132 + hidden_states = layer(positions, hidden_states)

133 return hidden_states

134

135

@@ -126,6 +139,7 @@ class BertLayer(nn.Module):

139 config: BertConfig,

140 cache_config: Optional[CacheConfig] = None,

141 quant_config: Optional[QuantizationConfig] = None,

142 + rotary_kwargs: Optional[dict] = None,

143 prefix: str = ""):

144 super().__init__()

145

@@ -135,6 +149,7 @@ class BertLayer(nn.Module):

149 layer_norm_eps=config.layer_norm_eps,

150 cache_config=cache_config,

151 quant_config=quant_config,

152 + rotary_kwargs=rotary_kwargs,

153 prefix=f"{prefix}.attention")

154

155 self.intermediate = BertIntermediate(

@@ -150,8 +165,8 @@ class BertLayer(nn.Module):

165 quant_config=quant_config,

166 prefix=f"{prefix}.output")

167

153 - def forward(self, hidden_states: torch.Tensor):

154 - attn_output = self.attention(hidden_states)

168 + def forward(self, positions: torch.Tensor, hidden_states: torch.Tensor):

169 + attn_output = self.attention(positions, hidden_states)

170 intermediate_output = self.intermediate(attn_output)

171 output = self.output(intermediate_output, attn_output)

172 return output

@@ -166,6 +181,7 @@ class BertAttention(nn.Module):

181 layer_norm_eps: float,

182 cache_config: Optional[CacheConfig] = None,

183 quant_config: Optional[QuantizationConfig] = None,

184 + rotary_kwargs: Optional[dict] = None,

185 prefix: str = "",

186 ):

187 super().__init__()

@@ -174,6 +190,7 @@ class BertAttention(nn.Module):

190 num_attention_heads=num_attention_heads,

191 cache_config=cache_config,

192 quant_config=quant_config,

193 + rotary_kwargs=rotary_kwargs,

194 prefix=f"{prefix}.output")

195

196 self.output = BertSelfOutput(hidden_size=hidden_size,

@@ -183,9 +200,10 @@ class BertAttention(nn.Module):

200

201 def forward(

202 self,

203 + positions: torch.Tensor,

204 hidden_states: torch.Tensor,

205 ) -> torch.Tensor:

188 - self_output = self.self(hidden_states)

206 + self_output = self.self(positions, hidden_states)

207 return self.output(self_output, hidden_states)

208

209

@@ -197,6 +215,7 @@ class BertSelfAttention(nn.Module):

215 num_attention_heads: int,

216 cache_config: Optional[CacheConfig] = None,

217 quant_config: Optional[QuantizationConfig] = None,

218 + rotary_kwargs: Optional[dict] = None,

219 prefix: str = "",

220 ):

221 super().__init__()

@@ -225,6 +244,11 @@ class BertSelfAttention(nn.Module):

244 quant_config=quant_config,

245 prefix=f"{prefix}.qkv_proj")

246

247 + if rotary_kwargs:

248 + self.rotary_emb = get_rope(**rotary_kwargs)

249 + else:

250 + self.rotary_emb = None

251 +

252 self.attn = Attention(num_heads=self.num_heads,

253 head_size=self.head_dim,

254 scale=self.scaling,

@@ -236,10 +260,15 @@ class BertSelfAttention(nn.Module):

260

261 def forward(

262 self,

263 + positions: torch.Tensor,

264 hidden_states: torch.Tensor,

265 ) -> torch.Tensor:

266 qkv, _ = self.qkv_proj(hidden_states)

267 q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

268 +

269 + if self.rotary_emb:

270 + q, k = self.rotary_emb(positions, q, k)

271 +

272 output = self.attn(q, k, v)

273 return output

274

@@ -321,11 +350,13 @@ class BertModel(nn.Module, SupportsQuant):

350 vllm_config: VllmConfig,

351 prefix: str = "",

352 embedding_class: type = BertEmbedding,

353 + rotary_kwargs: Optional[dict] = None,

354 add_pooling_layer: bool = False):

355 super().__init__()

356 config = vllm_config.model_config.hf_config

357 self.embeddings = embedding_class(config)

358 self.encoder = BertEncoder(vllm_config=vllm_config,

359 + rotary_kwargs=rotary_kwargs,

360 prefix=f"{prefix}.encoder")

361 self.pooler = BertPooler(config) if add_pooling_layer else None

362

@@ -347,7 +378,7 @@ class BertModel(nn.Module, SupportsQuant):

378 seq_lens=attn_metadata.seq_lens_tensor,

379 position_ids=position_ids,

380 token_type_ids=token_type_ids)

350 - return self.encoder(hidden_states)

381 + return self.encoder(position_ids, hidden_states)

382

383 def load_weights(self, weights: Iterable[Tuple[str,

384 torch.Tensor]]) -> Set[str]:

@@ -401,6 +432,7 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):

432 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

433 super().__init__()

434 pooler_config = vllm_config.model_config.pooler_config

435 + self.config = vllm_config.model_config.hf_config

436 self.model = self._build_model(vllm_config=vllm_config,

437 prefix=maybe_prefix(prefix, "model"))

438 self._pooler = self._build_pooler(pooler_config)

vllm/model_executor/models/roberta.py

+129 -64

@@ -22,30 +22,6 @@ from vllm.transformers_utils.config import (

22 from .interfaces import SupportsCrossEncoding, SupportsV0Only

25 -def roberta_task_weights_filter(

26 - all_weights: Iterable[Tuple[str, torch.Tensor]]

27 -) -> Tuple[Iterable[Tuple[str, torch.Tensor]], Iterable[Tuple[str,

28 - torch.Tensor]]]:

29 - """

30 - Separate task-specific weights that are applied on top

31 - of the encoder-decoder bert base.

32 - To do so, return two generators over the original iterator.

33 - Also, remove the "roberta." prefix to make it loadable

34 - from vanilla BertModel.

35 - """

36 - # Copy of a lazy iterator without in-memory overhead so both

37 - # iterators can be iterated upon independently.

38 - all_weights1, all_weights2 = itertools.tee(all_weights)

39 -

40 - def encoder_decoder_weights():

41 - for name, weight in all_weights1:

42 - if name.startswith("roberta."):

43 - yield (name[len("roberta."):], weight)

44 -

45 - return encoder_decoder_weights(), ((n, w) for n, w in all_weights2

46 - if not n.startswith("roberta."))

47 -

48 -

25 class RobertaEmbedding(nn.Module):

27 def __init__(self, config: RobertaConfig):

@@ -119,30 +95,6 @@ class RobertaEmbedding(nn.Module):

95 return embeddings

122 -# Adapted from transformers

123 -def create_position_ids_from_input_ids(input_ids,

124 - padding_idx,

125 - past_key_values_length=0):

126 - """

127 - Replace non-padding symbols with their position numbers.

128 - Position numbers begin at padding_idx+1. Padding symbols

129 - are ignored. This is modified from fairseq's `utils.make_positions`.

130 -

131 - Args:

132 - x: torch.Tensor x:

133 -

134 - Returns: torch.Tensor

135 - """

136 - # The series of casts and type-conversions here are carefully

137 - # balanced to both work with ONNX export and XLA.

138 - mask = input_ids.ne(padding_idx).int()

139 -

140 - incremental_indices = (torch.cumsum(mask, dim=0).type_as(mask) +

141 - past_key_values_length) * mask

142 -

143 - return incremental_indices.long() + padding_idx

144 -

145 -

98 # Adapted from transformers

99 class RobertaClassificationHead(nn.Module):

100 """Head for sentence-level classification tasks."""

@@ -174,15 +126,38 @@ class RobertaEmbeddingModel(BertEmbeddingModel):

126 def _build_model(self,

127 vllm_config: VllmConfig,

128 prefix: str = "") -> BertModel:

177 - return BertModel(vllm_config=vllm_config,

178 - prefix=prefix,

179 - embedding_class=RobertaEmbedding)

129 + if (vllm_config.model_config.hf_config.position_embedding_type ==

130 + "rotary"):

131 + config = vllm_config.model_config.hf_config

132 + head_dim = config.hidden_size // config.num_attention_heads

133 +

134 + rotary_kwargs = {

135 + "head_size": head_dim,

136 + "rotary_dim": getattr(config, "rotary_emb_dim", head_dim),

137 + "max_position": config.max_position_embeddings,

138 + "base": config.rotary_emb_base,

139 + "rope_scaling": getattr(config, "rope_scaling", None)

140 + }

141 +

142 + return BertModel(vllm_config=vllm_config,

143 + rotary_kwargs=rotary_kwargs,

144 + prefix=prefix)

145 + else:

146 + return BertModel(vllm_config=vllm_config,

147 + prefix=prefix,

148 + embedding_class=RobertaEmbedding)

149

150 def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):

151 + if getattr(self.config, "lora_rank", 0) > 0:

152 + scaling = self.config.lora_alpha / self.config.lora_rank

153 + weights = jina_merge_lora_weights(weights, scaling)

154 +

155 weights = self.hf_to_vllm_mapper.apply(weights)

156 # Separate weights in "roberta"-prefixed and all else (not in memory).

157 # For use with models like FacebookAI/roberta-base.

158 bert_weights, task_weights = roberta_task_weights_filter(weights)

159 + bert_weights = jina_to_vllm_mapper.apply(bert_weights)

160 +

161 loaded = self.model.load_weights(bert_weights)

162 if not len(loaded):

163 # Fix for models like `sentence-transformers/stsb-roberta-base-v2`

@@ -203,18 +178,6 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,

178 _pooler: An instance of Pooler used for pooling operations.

179 """

180

206 - jina_to_vllm_mapper = WeightsMapper(

207 - orig_to_new_substr={

208 - 'emb_ln': "embeddings.LayerNorm",

209 - 'layers': "layer",

210 - 'mixer.Wqkv': "attention.self.qkv_proj",

211 - 'mixer.out_proj': "attention.output.dense",

212 - 'norm1': "attention.output.LayerNorm",

213 - 'mlp.fc1': "intermediate.dense",

214 - 'mlp.fc2': "output.dense",

215 - 'norm2': "output.LayerNorm",

216 - })

217 -

181 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

182 super().__init__()

183 config = vllm_config.model_config.hf_config

@@ -232,7 +195,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,

195

196 def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):

197 bert_weights, task_weights = roberta_task_weights_filter(weights)

235 - bert_weights = self.jina_to_vllm_mapper.apply(bert_weights)

198 + bert_weights = jina_to_vllm_mapper.apply(bert_weights)

199

200 self.roberta.load_weights(bert_weights)

201

@@ -265,3 +228,105 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,

228 inputs_embeds=inputs_embeds,

229 intermediate_tensors=intermediate_tensors,

230 token_type_ids=token_type_ids)

231 +

232 +

233 +# Adapted from transformers

234 +def create_position_ids_from_input_ids(input_ids,

235 + padding_idx,

236 + past_key_values_length=0):

237 + """

238 + Replace non-padding symbols with their position numbers.

239 + Position numbers begin at padding_idx+1. Padding symbols

240 + are ignored. This is modified from fairseq's `utils.make_positions`.

241 +

242 + Args:

243 + x: torch.Tensor x:

244 +

245 + Returns: torch.Tensor

246 + """

247 + # The series of casts and type-conversions here are carefully

248 + # balanced to both work with ONNX export and XLA.

249 + mask = input_ids.ne(padding_idx).int()

250 +

251 + incremental_indices = (torch.cumsum(mask, dim=0).type_as(mask) +

252 + past_key_values_length) * mask

253 +

254 + return incremental_indices.long() + padding_idx

255 +

256 +

257 +def roberta_task_weights_filter(

258 + all_weights: Iterable[Tuple[str, torch.Tensor]]

259 +) -> Tuple[Iterable[Tuple[str, torch.Tensor]], Iterable[Tuple[str,

260 + torch.Tensor]]]:

261 + """

262 + Separate task-specific weights that are applied on top

263 + of the encoder-decoder bert base.

264 + To do so, return two generators over the original iterator.

265 + Also, remove the "roberta." prefix to make it loadable

266 + from vanilla BertModel.

267 + """

268 + # Copy of a lazy iterator without in-memory overhead so both

269 + # iterators can be iterated upon independently.

270 + all_weights1, all_weights2 = itertools.tee(all_weights)

271 +

272 + def encoder_decoder_weights():

273 + for name, weight in all_weights1:

274 + if name.startswith("roberta."):

275 + yield (name[len("roberta."):], weight)

276 +

277 + return encoder_decoder_weights(), ((n, w) for n, w in all_weights2

278 + if not n.startswith("roberta."))

279 +

280 +

281 +jina_to_vllm_mapper = WeightsMapper(

282 + orig_to_new_substr={

283 + 'emb_ln': "embeddings.LayerNorm",

284 + 'layers': "layer",

285 + 'mixer.Wqkv': "attention.self.qkv_proj",

286 + 'mixer.out_proj': "attention.output.dense",

287 + 'norm1': "attention.output.LayerNorm",

288 + 'mlp.fc1': "intermediate.dense",

289 + 'mlp.fc2': "output.dense",

290 + 'norm2': "output.LayerNorm",

291 + })

292 +

293 +

294 +@torch.inference_mode()

295 +def jina_merge_lora_weights(weights: Iterable[Tuple[str, torch.Tensor]],

296 + scaling: float = 1.0):

297 + # use for jina-embeddings-v3

298 + # Merge Lora weights into a single weight tensor.

299 + # This is a temporary solution until we have a better way to handle

300 +

301 + weights = {name: weight for name, weight in weights}

302 +

303 + o = ".original"

304 + a = ".0.lora_A"

305 + b = ".0.lora_B"

306 +

307 + # text-matching

308 + i = -1

309 +

310 + for name in list(weights.keys()):

311 + if o in name:

312 + dtype = weights[name].dtype

313 + shape = weights[name].shape

314 + weight_name = name[:-len(o)]

315 +

316 + if "embeddings" in weight_name:

317 + B = weights[weight_name + a][i].cuda().float()

318 + A = weights[weight_name + b][i].cuda().float()

319 + else:

320 + B = weights[weight_name + b][i].cuda().float()

321 + A = weights[weight_name + a][i].cuda().float()

322 +

323 + weight = (weights[weight_name + o].cuda() +

324 + torch.matmul(B, A).view(shape) * scaling)

325 + weight = weight.cpu().to(dtype)

326 +

327 + weights[weight_name.replace(".parametrizations", "")] = weight

328 +

329 + del weights[weight_name + o], weights[weight_name +

330 + a], weights[weight_name + b]

331 +

332 + return [(name, weight) for name, weight in weights.items()]