File size: 1,012 Bytes
1b7e88c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
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
|