Fix bug for MaskedLanguageModel class` (#513)

Signed-off-by: Liu, Kaixuan <kaixuan.liu@intel.com>

Author: kaixuanliu

backends/python/server/text_embeddings_server/models/__init__.py

@@ -54,12 +54,11 @@ def get_model(model_path: Path, dtype: Optional[str], pool: str):
             and FLASH_ATTENTION
         ):
             if pool != "cls":
-                if config.architectures[0].endswith("ForMaskedLM"):
+                if config.architectures[0].endswith("ForMaskedLM") and pool == "splade":
                     return MaskedLanguageModel(
                         model_path,
                         device,
                         datatype,
-                        pool,
                         trust_remote=TRUST_REMOTE_CODE,
                     )
                 return DefaultModel(
@@ -70,9 +69,9 @@ def get_model(model_path: Path, dtype: Optional[str], pool: str):
             return ClassificationModel(
                 model_path, device, datatype, trust_remote=TRUST_REMOTE_CODE
             )
-        elif config.architectures[0].endswith("ForMaskedLM"):
+        elif config.architectures[0].endswith("ForMaskedLM") and pool == "splade":
             return MaskedLanguageModel(
-                model_path, device, datatype, pool, trust_remote=TRUST_REMOTE_CODE
+                model_path, device, datatype, trust_remote=TRUST_REMOTE_CODE
             )
         else:
             return DefaultModel(
@@ -97,9 +96,9 @@ def get_model(model_path: Path, dtype: Optional[str], pool: str):
                     datatype,
                     trust_remote=TRUST_REMOTE_CODE,
                 )
-            elif config.architectures[0].endswith("ForMaskedLM"):
-                return MaskedLanguageModel(
-                    model_path, device, datatype, pool, trust_remote=TRUST_REMOTE_CODE
+            elif config.architectures[0].endswith("ForMaskedLM") and pool == "splade":
+                model_handle = MaskedLanguageModel(
+                    model_path, device, datatype, trust_remote=TRUST_REMOTE_CODE
                 )
             else:
                 model_handle = DefaultModel(
@@ -119,9 +118,9 @@ def get_model(model_path: Path, dtype: Optional[str], pool: str):
                     datatype,
                     trust_remote=TRUST_REMOTE_CODE,
                 )
-            elif config.architectures[0].endswith("ForMaskedLM"):
+            elif config.architectures[0].endswith("ForMaskedLM") and pool == "splade":
                 return MaskedLanguageModel(
-                    model_path, device, datatype, pool, trust_remote=TRUST_REMOTE_CODE
+                    model_path, device, datatype, trust_remote=TRUST_REMOTE_CODE
                 )
             else:
                 return DefaultModel(

backends/python/server/text_embeddings_server/models/default_model.py

@@ -5,7 +5,7 @@
 from typing import Type, List
 from transformers import AutoModel
 from opentelemetry import trace
-from sentence_transformers.models import Pooling
+from text_embeddings_server.models.pooling import DefaultPooling
 
 from text_embeddings_server.models import Model
 from text_embeddings_server.models.types import PaddedBatch, Embedding, Score
@@ -28,7 +28,7 @@ def __init__(
             .to(device)
         )
         self.hidden_size = model.config.hidden_size
-        self.pooling = Pooling(self.hidden_size, pooling_mode=pool)
+        self.pooling = DefaultPooling(self.hidden_size, pooling_mode=pool)
 
         position_offset = 0
         model_type = model.config.model_type
@@ -65,11 +65,7 @@ def embed(self, batch: PaddedBatch) -> List[Embedding]:
             kwargs["position_ids"] = batch.position_ids
         output = self.model(**kwargs)
 
-        pooling_features = {
-            "token_embeddings": output[0],
-            "attention_mask": batch.attention_mask,
-        }
-        embedding = self.pooling.forward(pooling_features)["sentence_embedding"]
+        embedding = self.pooling.forward(output, batch.attention_mask)
 
         cpu_results = embedding.view(-1).tolist()
 

backends/python/server/text_embeddings_server/models/masked_model.py

@@ -8,7 +8,7 @@
 
 from text_embeddings_server.models import Model
 from text_embeddings_server.models.types import PaddedBatch, Embedding, Score
-from text_embeddings_server.models.pooling import DefaultPooling, SpladePooling
+from text_embeddings_server.models.pooling import SpladePooling
 
 tracer = trace.get_tracer(__name__)
 
@@ -19,7 +19,6 @@ def __init__(
         model_path: Path,
         device: torch.device,
         dtype: torch.dtype,
-        pool: str,
         trust_remote: bool = False,
     ):
         model = (
@@ -29,14 +28,17 @@ def __init__(
             .to(dtype)
             .to(device)
         )
-        self.hidden_size = model.config.hidden_size
-        self.vocab_size = model.config.vocab_size
-        self.pooling_mode = pool
-        if pool == "splade":
-            self.pooling = SpladePooling()
+        self.pooling = SpladePooling()
+        position_offset = 0
+        model_type = model.config.model_type
+        if model_type in ["xlm-roberta", "camembert", "roberta"]:
+            position_offset = model.config.pad_token_id + 1
+        if hasattr(model.config, "max_seq_length"):
+            self.max_input_length = model.config.max_seq_length
         else:
-            self.pooling = DefaultPooling(self.hidden_size, pooling_mode=pool)
-
+            self.max_input_length = (
+                model.config.max_position_embeddings - position_offset
+            )
         self.has_position_ids = (
             inspect.signature(model.forward).parameters.get("position_ids", None)
             is not None

backends/python/server/text_embeddings_server/models/types.py

@@ -1,4 +1,5 @@
 import os
+import math
 import torch
 
 from abc import ABC, abstractmethod