omagent_core/models/encoders/base.py

from abc import ABC, abstractmethod
from typing import List

from ...base import BotBase
from ...utils.general import chunks


class EncoderBase(BotBase, ABC):
    endpoint: str
    dim: int  # Dimension of the vector
    batch_size: int = 20

    @abstractmethod
    def _infer(self, data: List[str], **kwargs) -> List[List[float]]:
        """Encoding"""

    async def _ainfer(self, data: List[str], **kwargs) -> List[List[float]]:
        """Async encoding"""
        raise NotImplementedError("Async generation not implemented for this Encoder.")

    def infer(self, data: List[str], **kwargs) -> List[List[float]]:
        res = []
        for chunk in chunks(data, self.batch_size, self.batch_size):
            res += self._infer(chunk, **kwargs)
        return res

    async def ainfer(self, data: List[str], **kwargs) -> List[List[float]]:
        res = []
        for chunk in chunks(data, self.batch_size, self.batch_size):
            res += await self._ainfer(chunk, **kwargs)
        return res