Fixed embeddings

predibase · Oct 18, 2024 · df5ae30 · df5ae30
1 parent fcfa679
commit df5ae30
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 14 deletions.
diff --git a/server/lorax_server/models/flash_causal_lm.py b/server/lorax_server/models/flash_causal_lm.py
@@ -353,19 +353,6 @@ def from_pb(
             adapter_meta=None,
         )
 
-    @classmethod
-    def from_pb_embed(
-        self,
-        pb: generate_pb2.EmbedRequest,
-        tokenizer: PreTrainedTokenizerBase,
-        tokenizers: TokenizerManager,
-        processor,
-        config,
-        dtype,
-        device,
-    ) -> "FlashCausalLMBatch":
-        return self.from_pb(pb, tokenizer, tokenizers, None, None, dtype, device)
-
     @tracer.start_as_current_span("filter")
     def filter(self, request_ids: List[int]) -> "FlashCausalLMBatch":
         if len(request_ids) == 0:

diff --git a/server/lorax_server/server.py b/server/lorax_server/server.py
@@ -138,7 +138,7 @@ async def Embed(self, request: generate_pb2.EmbedRequest, context):
         if not self.model.supports_embeddings:
             raise ValueError("Model does not support embeddings")
 
-        batch = self.model.batch_type.from_pb_embed(
+        batch = self.model.batch_type.from_pb(
             request.batch,
             self.model.tokenizer,
             self.model.tokenizers,