Fix embedding generation for T5 models

app.py

@@ -43,8 +43,15 @@ def startup_event():
     if transformers_direct_tokenize is not None and transformers_direct_tokenize == "true" or transformers_direct_tokenize == "1":
         direct_tokenize = True
 
-    meta_config = Meta('./models/model')
-    vec = Vectorizer('./models/model', cuda_support, cuda_core, cuda_per_process_memory_fraction,
+    def get_model_directory() -> str:
+        if os.path.exists("./models/model/model_name"):
+            with open("./models/model/model_name", "r") as f:
+                model_name = f.read()
+                return f"./models/model/{model_name}"
+        return "./models/model"
+
+    meta_config = Meta(get_model_directory())
+    vec = Vectorizer(get_model_directory(), cuda_support, cuda_core, cuda_per_process_memory_fraction,
                      meta_config.getModelType(), meta_config.get_architecture(), direct_tokenize)
 
 

download.py

@@ -8,8 +8,10 @@
     AutoTokenizer,
     AutoConfig,
 )
+from sentence_transformers import SentenceTransformer
 
 
+model_dir = './models/model'
 model_name = os.getenv('MODEL_NAME', None)
 force_automodel = os.getenv('FORCE_AUTOMODEL', False)
 if not model_name:
@@ -22,21 +24,28 @@
 
 print(f"Downloading model {model_name} from huggingface model hub")
 config = AutoConfig.from_pretrained(model_name)
-if config.architectures and not force_automodel:
-    print(f"Using class {config.architectures[0]} to load model weights")
-    mod = __import__('transformers', fromlist=[config.architectures[0]])
-    try:
-        klass_architecture = getattr(mod, config.architectures[0])
-        model = klass_architecture.from_pretrained(model_name)
-    except AttributeError:
-        print(f"{config.architectures[0]} not found in transformers, fallback to AutoModel")
-        model = AutoModel.from_pretrained(model_name)
+model_type = config.to_dict()['model_type']
+
+if model_type is not None and model_type == "t5":
+    SentenceTransformer(model_name, cache_folder=model_dir)
+    with open(f"{model_dir}/model_name", "w") as f:
+        f.write(model_name.replace("/", "_"))
 else:
-    model = AutoModel.from_pretrained(model_name)
+    if config.architectures and not force_automodel:
+        print(f"Using class {config.architectures[0]} to load model weights")
+        mod = __import__('transformers', fromlist=[config.architectures[0]])
+        try:
+            klass_architecture = getattr(mod, config.architectures[0])
+            model = klass_architecture.from_pretrained(model_name)
+        except AttributeError:
+            print(f"{config.architectures[0]} not found in transformers, fallback to AutoModel")
+            model = AutoModel.from_pretrained(model_name)
+    else:
+        model = AutoModel.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-model.save_pretrained('./models/model')
-tokenizer.save_pretrained('./models/model')
+    model.save_pretrained(model_dir)
+    tokenizer.save_pretrained(model_dir)
 
 nltk.download('punkt', download_dir='./nltk_data')

requirements-test.txt

@@ -5,4 +5,5 @@ uvicorn==0.21.1
 nltk==3.8.1
 torch==2.0.0
 sentencepiece==0.1.97
+sentence-transformers==2.2.2
 pytest

requirements.txt

@@ -1,6 +1,7 @@
-transformers==4.29.2
-fastapi==0.95.2
-uvicorn==0.22.0
+transformers==4.34.0
+fastapi==0.103.2
+uvicorn==0.23.2
 nltk==3.8.1
-torch==2.0.1
+torch==2.1.0
 sentencepiece==0.1.99
+sentence-transformers==2.2.2

smoke_test.py

@@ -50,6 +50,7 @@ def try_to_vectorize(url):
             # aware of 384 and 768 dim vectors, which should both fall in that
             # range
             self.assertTrue(len(resBody['vector']) > 100)
+            print(f"vector dimensions are: {len(resBody['vector'])}")
 
         try_to_vectorize(self.url + "/vectors/")
         try_to_vectorize(self.url + "/vectors")

vectorizer.py

@@ -14,6 +14,7 @@
     DPRContextEncoder,
     DPRQuestionEncoder,
 )
+from sentence_transformers import SentenceTransformer
 
 
 # limit transformer batch size to limit parallel inference, otherwise we run
@@ -29,14 +30,47 @@ class VectorInput(BaseModel):
     text: str
     config: Optional[VectorInputConfig] = None
 
+
 class Vectorizer:
+    executor: ThreadPoolExecutor
+
+    def __init__(self, model_path: str, cuda_support: bool, cuda_core: str, cuda_per_process_memory_fraction: float, model_type: str, architecture: str, direct_tokenize: bool):
+        self.executor = ThreadPoolExecutor()
+        if model_type == 't5':
+            self.vectorizer = SentenceTransformerVectorizer(model_path, cuda_core)
+        else:
+            self.vectorizer = HuggingFaceVectorizer(model_path, cuda_support, cuda_core, cuda_per_process_memory_fraction, model_type, architecture, direct_tokenize)
+
+    async def vectorize(self, text: str, config: VectorInputConfig):
+        return await asyncio.wrap_future(self.executor.submit(self.vectorizer.vectorize, text, config))
+
+
+class SentenceTransformerVectorizer:
+    model: SentenceTransformer
+    cuda_core: str
+
+    def __init__(self, model_path: str, cuda_core: str):
+        self.cuda_core = cuda_core
+        self.model = SentenceTransformer(model_path, device=self.get_device())
+        self.model.eval() # make sure we're in inference mode, not training
+
+    def get_device(self) -> Optional[str]:
+        if self.cuda_core is not None and self.cuda_core != "":
+            return self.cuda_core
+        return None
+
+    def vectorize(self, text: str, config: VectorInputConfig):
+        embedding = self.model.encode([text], device=self.get_device(), convert_to_tensor=False, convert_to_numpy=True)
+        return embedding[0]
+
+
+class HuggingFaceVectorizer:
     model: AutoModel
     tokenizer: AutoTokenizer
     cuda: bool
     cuda_core: str
     model_type: str
     direct_tokenize: bool
-    executor: ThreadPoolExecutor
 
     def __init__(self, model_path: str, cuda_support: bool, cuda_core: str, cuda_per_process_memory_fraction: float, model_type: str, architecture: str, direct_tokenize: bool):
         self.cuda = cuda_support
@@ -56,8 +90,6 @@ def __init__(self, model_path: str, cuda_support: bool, cuda_core: str, cuda_per
 
         self.tokenizer = self.model_delegate.create_tokenizer(model_path)
 
-        self.executor = ThreadPoolExecutor()
-
         nltk.data.path.append('./nltk_data')
 
     def tokenize(self, text:str):
@@ -73,7 +105,7 @@ def get_batch_results(self, tokens, text):
     def pool_embedding(self, batch_results, tokens, config):
         return self.model_delegate.pool_embedding(batch_results, tokens, config)
 
-    def _vectorize(self, text: str, config: VectorInputConfig):
+    def vectorize(self, text: str, config: VectorInputConfig):
         with torch.no_grad():
             if self.direct_tokenize:
                 # create embeddings without tokenizing text
@@ -100,9 +132,6 @@ def _vectorize(self, text: str, config: VectorInputConfig):
                     batch_sum_vectors += self.pool_embedding(batch_results, tokens, config)
                 return batch_sum_vectors.detach() / num_sentences
 
-    async def vectorize(self, text: str, config: VectorInputConfig):
-        return await asyncio.wrap_future(self.executor.submit(self._vectorize, text, config))
-
 
 class HFModel: