INNER CODE UNIT · Python

model_output

tensorchord/modelz-llm · src/modelz_llm/emb.py:46

        model_output = self.model(**inputs)
        # Perform pooling
        sentence_embeddings = mean_pooling(model_output, inputs["attention_mask"])
        # Normalize embeddings
        sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)

        return token_count, sentence_embeddings

    def __call__(self, req: EmbeddingRequest) -> EmbeddingResponse:
        token_count, embeddings = self.get_embedding_with_token_count(req.input)
        embeddings = embeddings.detach()
        if self.device != "cpu":
            embeddings = embeddings.cpu()
        embeddings = embeddings.numpy()
        if req.encoding_format == "base64":
            embeddings = [
                base64.b64encode(emb.astype(np.float32).tobytes()).decode("utf-8")
                for emb in embeddings

View source record →