| 12345678910111213141516171819202122232425262728293031 |
- from typing import Any
- from PIL.Image import Image
- from sentence_transformers import SentenceTransformer
- from sentence_transformers.util import snapshot_download
- from ..schemas import ModelType
- from .base import InferenceModel
- class CLIPSTEncoder(InferenceModel):
- _model_type = ModelType.CLIP
- def _download(self, **model_kwargs: Any) -> None:
- repo_id = self.model_name if "/" in self.model_name else f"sentence-transformers/{self.model_name}"
- snapshot_download(
- cache_dir=self.cache_dir,
- repo_id=repo_id,
- library_name="sentence-transformers",
- ignore_files=["flax_model.msgpack", "rust_model.ot", "tf_model.h5"],
- )
- def _load(self, **model_kwargs: Any) -> None:
- self.model = SentenceTransformer(
- self.model_name,
- cache_folder=self.cache_dir.as_posix(),
- **model_kwargs,
- )
- def _predict(self, image_or_text: Image | str) -> list[float]:
- return self.model.encode(image_or_text).tolist()
|
