diff --git a/semantic_router/encoders/fastembed.py b/semantic_router/encoders/fastembed.py index c730d4c68d59ca5c6cabbcfa2ba9be9bfe79ef6a..fb845ce7278470f790b765d19cba0c318e04c6fb 100644 --- a/semantic_router/encoders/fastembed.py +++ b/semantic_router/encoders/fastembed.py @@ -1,12 +1,14 @@ from typing import Any, List, Optional import numpy as np -from pydantic import BaseModel, PrivateAttr +from pydantic import PrivateAttr +from semantic_router.encoders import BaseEncoder -class FastEmbedEncoder(BaseModel): + +class FastEmbedEncoder(BaseEncoder): type: str = "fastembed" - model_name: str = "BAAI/bge-small-en-v1.5" + name: str = "BAAI/bge-small-en-v1.5" max_length: int = 512 cache_dir: Optional[str] = None threads: Optional[int] = None @@ -27,7 +29,7 @@ class FastEmbedEncoder(BaseModel): ) embedding_args = { - "model_name": self.model_name, + "model_name": self.name, "max_length": self.max_length, "cache_dir": self.cache_dir, "threads": self.threads, diff --git a/semantic_router/schema.py b/semantic_router/schema.py index 826cb6864d561c98cd54b7eac9c4d2eb592777ab..644803556176fbd5f7c70ee4040b686c9207c38e 100644 --- a/semantic_router/schema.py +++ b/semantic_router/schema.py @@ -36,11 +36,11 @@ class Encoder: if self.type == EncoderType.HUGGINGFACE: raise NotImplementedError elif self.type == EncoderType.FASTEMBED: - self.model = FastEmbedEncoder(name) + self.model = FastEmbedEncoder(name=name) elif self.type == EncoderType.OPENAI: - self.model = OpenAIEncoder(name) + self.model = OpenAIEncoder(name=name) elif self.type == EncoderType.COHERE: - self.model = CohereEncoder(name) + self.model = CohereEncoder(name=name) else: raise ValueError