tei support

config.json

@@ -1,51 +1,20 @@
 {
-  "_name_or_path": "jina-embeddings-v3-xlm-large-pairs-vol5-lora",
-  "architectures": [
-    "XLMRobertaModel"
-  ],
-  "attention_probs_dropout_prob": 0.1,
-  "auto_map": {
-    "AutoConfig": "jinaai/xlm-roberta-flash-implementation--configuration_xlm_roberta.XLMRobertaFlashConfig",
-    "AutoModel": "jinaai/xlm-roberta-flash-implementation--modeling_lora.XLMRobertaLoRA",
-    "AutoModelForMaskedLM": "jinaai/xlm-roberta-flash-implementation--modeling_xlm_roberta.XLMRobertaForMaskedLM",
-    "AutoModelForPreTraining": "jinaai/xlm-roberta-flash-implementation--modeling_xlm_roberta.XLMRobertaForPreTraining"
-  },
-  "bos_token_id": 0,
-  "classifier_dropout": null,
-  "emb_pooler": null,
-  "eos_token_id": 2,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.1,
+  "model_type": "new",
+  "vocab_size": 250002,
   "hidden_size": 1024,
-  "initializer_range": 0.02,
+  "num_hidden_layers": 24,
+  "num_attention_heads": 16,
   "intermediate_size": 4096,
-  "layer_norm_eps": 1e-05,
-  "load_trained_adapters": true,
-  "lora_adaptations": ["retrieval.query", "retrieval.passage", "separation", "classification", "text-matching"],
-  "lora_alpha": 1,
-  "lora_dropout_p": 0.0,
-  "lora_main_params_trainable": false,
-  "lora_rank": 4,
-  "matryoshka_dimensions": [32, 64, 128, 256, 512, 768, 1024],
+  "hidden_act": "gelu",
   "max_position_embeddings": 8194,
-  "num_attention_heads": 16,
-  "num_hidden_layers": 24,
-  "output_past": true,
-  "pad_token_id": 1,
-  "position_embedding_type": "rotary",
-  "rotary_emb_base": 20000.0,
-  "torch_dtype": "bfloat16",
-  "transformers_version": "4.30.2",
-  "truncate_dim": null,
   "type_vocab_size": 1,
-  "use_cache": true,
-  "use_flash_attn": true,
-  "vocab_size": 250002,
-  "task_instructions": {
-    "retrieval.query": "Represent the query for retrieving evidence documents: ",
-    "retrieval.passage": "Represent the document for retrieval: ",
-    "separation": "",
-    "classification": "",
-    "text-matching": ""
-  }
+  "layer_norm_type": "layer_norm",
+  "layer_norm_eps": 1e-5,
+  "position_embedding_type": "rope",
+  "rope_theta": 10000.0,
+  "rope_scaling": null,
+  "logn_attention_scale": false,
+  "logn_attention_clip1": false,
+  "architectures": ["GTEModel"],
+  "model_max_length": 8194
 }

conv.py

@@ -0,0 +1,71 @@
+from safetensors import safe_open
+from safetensors.torch import save_file
+import torch
+
+def rename_key(key):
+    parts = key.split('.')
+    if 'roberta' in parts:
+        parts.remove('roberta')
+    if 'parametrizations' in parts:
+        parts.remove('parametrizations')
+    if 'weight' in parts and 'original' in parts:
+        parts.remove('original')
+    
+    if 'encoder.layers' in key:
+        parts[parts.index('layers')] = 'layer'
+    
+    if 'mixer' in parts:
+        parts[parts.index('mixer')] = 'attention'
+    if 'Wqkv' in parts:
+        parts[parts.index('Wqkv')] = 'qkv_proj'
+    if 'out_proj' in parts:
+        parts[parts.index('out_proj')] = 'o_proj'
+    if 'mlp.fc1' in key:
+        parts[parts.index('fc1')] = 'up_proj'
+    if 'mlp.fc2' in key:
+        parts[parts.index('fc2')] = 'down_proj'
+    if 'emb_ln' in parts:
+        parts[parts.index('emb_ln')] = 'LayerNorm'
+        parts.insert(0, 'embeddings')
+    if 'norm1' in parts:
+        parts[parts.index('norm1')] = 'attn_ln'
+    if 'norm2' in parts:
+        parts[parts.index('norm2')] = 'mlp_ln'
+    if 'weight' in parts:
+        if parts[-2] in ['attn_ln', 'mlp_ln', 'LayerNorm']:
+            parts[-1] = 'gamma'
+    if 'bias' in parts:
+        if parts[-2] in ['attn_ln', 'mlp_ln', 'LayerNorm']:
+            parts[-1] = 'beta'
+    
+    return '.'.join(parts)
+
+input_file = "original_model.safetensors"
+output_file = "model.safetensors"
+
+new_tensors = {}
+
+with safe_open(input_file, framework="pt", device="cpu") as f:
+    for key in f.keys():
+        if 'lora' not in key:
+            new_key = rename_key(key)
+            tensor = f.get_tensor(key)
+            
+            if 'mlp.up_proj' in new_key:
+                # Create up_proj and up_gate_proj
+                new_tensors[new_key] = tensor
+                gate_key = new_key.replace('up_proj', 'up_gate_proj')
+                # Expand the tensor to match the expected shape
+                expanded_tensor = torch.cat([tensor] * 2, dim=0)
+                new_tensors[gate_key] = expanded_tensor
+            else:
+                new_tensors[new_key] = tensor
+
+save_file(new_tensors, output_file)
+print(f"Renamed tensors saved to {output_file}")
+
+# Inspect the renamed tensors
+with safe_open(output_file, framework="pt", device="cpu") as f:
+    print("\nRenamed tensors:")
+    for key in f.keys():
+        print(f"{key}: {f.get_tensor(key).shape}")

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:17ca06efd886a065d0081912b04c9e27ef5086a9dd09659cce32aa9c84587f23
-size 1144685320
+oid sha256:dd20adef5134ed345186af0acc5066662700c49c5c328a5d358f60f176f274a3
+size 1521812032

original_config.json

@@ -0,0 +1,51 @@
+{
+  "_name_or_path": "jina-embeddings-v3-xlm-large-pairs-vol5-lora",
+  "architectures": [
+    "XLMRobertaModel"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "jinaai/xlm-roberta-flash-implementation--configuration_xlm_roberta.XLMRobertaFlashConfig",
+    "AutoModel": "jinaai/xlm-roberta-flash-implementation--modeling_lora.XLMRobertaLoRA",
+    "AutoModelForMaskedLM": "jinaai/xlm-roberta-flash-implementation--modeling_xlm_roberta.XLMRobertaForMaskedLM",
+    "AutoModelForPreTraining": "jinaai/xlm-roberta-flash-implementation--modeling_xlm_roberta.XLMRobertaForPreTraining"
+  },
+  "bos_token_id": 0,
+  "classifier_dropout": null,
+  "emb_pooler": null,
+  "eos_token_id": 2,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "load_trained_adapters": true,
+  "lora_adaptations": ["retrieval.query", "retrieval.passage", "separation", "classification", "text-matching"],
+  "lora_alpha": 1,
+  "lora_dropout_p": 0.0,
+  "lora_main_params_trainable": false,
+  "lora_rank": 4,
+  "matryoshka_dimensions": [32, 64, 128, 256, 512, 768, 1024],
+  "max_position_embeddings": 8194,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "output_past": true,
+  "pad_token_id": 1,
+  "position_embedding_type": "rotary",
+  "rotary_emb_base": 20000.0,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.30.2",
+  "truncate_dim": null,
+  "type_vocab_size": 1,
+  "use_cache": true,
+  "use_flash_attn": true,
+  "vocab_size": 250002,
+  "task_instructions": {
+    "retrieval.query": "Represent the query for retrieving evidence documents: ",
+    "retrieval.passage": "Represent the document for retrieval: ",
+    "separation": "",
+    "classification": "",
+    "text-matching": ""
+  }
+}

original_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17ca06efd886a065d0081912b04c9e27ef5086a9dd09659cce32aa9c84587f23
+size 1144685320

resolve/main/config.json

@@ -0,0 +1,25 @@
+{
+  "_name_or_path": "jina-embeddings-v3-xlm-large-pairs-vol5-lora",
+  "architectures": [
+    "XLMRobertaModel"
+  ],
+  "model_type": "xlm-roberta",
+  "max_position_embeddings": 8194,
+  "pad_token_id": 1,
+  "attention_probs_dropout_prob": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-05,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "vocab_size": 250002,
+  "type_vocab_size": 1,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "classifier_dropout": null,
+  "output_past": true,
+  "position_embedding_type": "absolute"
+}