add llm input_embeddings layer to enable text-to-vector capabilities, and uploade script for extracting GPT-NeoX embedding layer.

Author: peng3307165

docs/script/get_input_embedding_script.py

@@ -0,0 +1,11 @@
+# -*- coding: utf-8 -*-
+import torch
+import numpy as np
+from transformers import AutoModelForCausalLM
+
+
+model_path = ''
+device = torch.device('cuda')
+model = AutoModelForCausalLM.from_pretrained(model_path, local_files_only=True).to(device)
+embedding_weights = model.get_input_embeddings().weight.to('cpu').detach().numpy()
+np.save('gpt-neox-embedding.npy', embedding_weights)

modelcache/embedding/__init__.py

@@ -2,6 +2,7 @@
 from modelcache.utils.lazy_import import LazyImport
 huggingface = LazyImport("huggingface", globals(), "modelcache.embedding.huggingface")
 data2vec = LazyImport("data2vec", globals(), "modelcache.embedding.data2vec")
+llmEmb = LazyImport("llmEmb", globals(), "modelcache.embedding.llmEmb")
 
 
 def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
@@ -10,3 +11,7 @@ def Huggingface(model="sentence-transformers/all-mpnet-base-v2"):
 
 def Data2VecAudio(model="facebook/data2vec-audio-base-960h"):
     return data2vec.Data2VecAudio(model)
+
+
+def LlmEmb2vecAudio():
+    return llmEmb.LlmEmb2Vec()

modelcache/embedding/llmEmb.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+from modelcache.embedding.base import BaseEmbedding
+from transformers import AutoTokenizer
+from transformers import AutoConfig
+
+
+class LlmEmb2Vec(BaseEmbedding):
+    def __init__(self):
+
+        self.model_name = ''  # 13b-mft-embedding.npy
+        model_path = ''  # .npy file storage path
+        model_file = model_path + self.model_name  # .npy file
+        config = AutoConfig.from_pretrained(model_path)
+        dimension = config.hidden_size
+        self.__dimension = dimension
+        self.model = np.load(model_file)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
+
+    def to_embeddings(self, data, **_):
+        """Generate embedding given text input
+
+        :param data: text in string.
+        :return: a text embedding in shape of (dim,).
+        """
+        input_ids = self.tokenizer.encode(data, add_special_tokens=True)
+        embedding_array = self.model[input_ids].mean(axis=0)
+        return embedding_array
+
+    def post_proc(self, token_embeddings, inputs):
+        pass
+
+    @property
+    def dimension(self):
+        """Embedding dimension.
+        :return: embedding dimension
+        """
+        return self.__dimension