fix modeing

Author: DrownFish19

paddlenlp/transformers/__init__.py

@@ -215,6 +215,7 @@
 from .layoutxlm.modeling import *
 from .layoutxlm.tokenizer import *
 from .llama import *
+from .llm_embed.modeling import *
 from .luke.configuration import *
 from .luke.modeling import *
 from .luke.tokenizer import *
@@ -246,6 +247,7 @@
 from .nezha.configuration import *
 from .nezha.modeling import *
 from .nezha.tokenizer import *
+from .nv_embed.modeling import *
 from .nystromformer.configuration import *
 from .nystromformer.modeling import *
 from .nystromformer.tokenizer import *
@@ -320,5 +322,3 @@
 from .xlm_roberta.tokenizer import *
 from .xlm_roberta.configuration import *
 from .yuan import *
-from .llm_embed import *
-from .nv_embed import *

paddlenlp/transformers/llm_embed/modeling.py

@@ -34,10 +34,14 @@ class EncoderOutput(ModelOutput):
     scores: Optional[paddle.Tensor] = None
 
 
+__all__ = ["BiEncoderModel"]
+
+
 class BiEncoderModel(PretrainedModel):
     def __init__(
         self,
         model_name_or_path: str = None,
+        dtype: str = "float16",
         normalized: bool = False,
         sentence_pooling_method: str = "cls",
         negatives_cross_device: bool = False,
@@ -53,7 +57,7 @@ def __init__(
         max_seq_length: int = 4096,
     ):
         super().__init__()
-        self.model = AutoModel.from_pretrained(model_name_or_path, convert_from_torch=True)
+        self.model = AutoModel.from_pretrained(model_name_or_path, dtype=dtype, convert_from_torch=True)
         self.model_config = AutoConfig.from_pretrained(model_name_or_path)
         self.cross_entropy = nn.CrossEntropyLoss(reduction="mean")
 

paddlenlp/transformers/nv_embed/modeling.py

@@ -26,6 +26,8 @@
 from .. import AutoTokenizer, MistralModel, PretrainedConfig, PretrainedModel
 from ..model_outputs import BaseModelOutputWithPast, ModelOutput
 
+__all__ = ["NVEncodeModel"]
+
 
 @dataclass
 class EncoderOutput(ModelOutput):
@@ -121,7 +123,7 @@ def forward(self, last_hidden_states, pool_mask):
         one = paddle.eye(
             num_rows=self.config.hidden_size,
             num_columns=self.config.hidden_size,
-            dtype=str(self.latents.weight.dtype).split(".")[-1],
+            dtype=self.latents.weight.dtype,
         )
         self_latents_weight_T = self.latents(one).T
         # latents = repeat(self_latents_weight_T, "d h -> b d h", b=last_hidden_states.shape[0]) # from einops import repeat
@@ -206,7 +208,8 @@ def __init__(
         self.latent_model = LatentModel(config=config)  # get latent model structure
 
         self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, padding_side="right")
-        self.tokenizer.pad_token = self.tokenizer.eos_token
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
 
         self.query_instruction = query_instruction
         self.document_instruction = document_instruction
@@ -495,6 +498,7 @@ def encode_sentences(self, sentences: List[str], instruction_len, **kwargs) -> n
                 sentences_batch,
                 max_length=4096,
                 padding=True,
+                return_attention_mask=True,
                 return_token_type_ids=False,
                 return_tensors="pd",
                 truncation=True,

tests/transformers/llm_embed/test_modeling.py

@@ -19,11 +19,10 @@
 
 from paddlenlp.transformers import AutoTokenizer, BiEncoderModel
 
-from ...testing_utils import require_gpu, slow
+from ...testing_utils import require_gpu
 
 
 class BiEncoderModelIntegrationTest(unittest.TestCase):
-    @slow
     @require_gpu(1)
     def test_model_tiny_logits(self):
         input_texts = [

tests/transformers/nv_embed/test_modeling.py

@@ -19,11 +19,10 @@
 
 from paddlenlp.transformers import NVEncodeModel, PretrainedConfig
 
-from ...testing_utils import require_gpu, slow
+from ...testing_utils import require_gpu
 
 
 class NVEncodeModelIntegrationTest(unittest.TestCase):
-    @slow
     @require_gpu(1)
     def test_model_tiny_logits(self):
         input_texts = [