Skip to content
Snippets Groups Projects
Unverified Commit 5006b49b authored by James Briggs's avatar James Briggs
Browse files

lint fixes

parent c9715a34
Branches
Tags v0.2.5
No related merge requests found
from typing import Any, List, Optional from typing import Any, List, Optional
import numpy as np 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" type: str = "fastembed"
model_name: str = "BAAI/bge-small-en-v1.5" name: str = "BAAI/bge-small-en-v1.5"
max_length: int = 512 max_length: int = 512
cache_dir: Optional[str] = None cache_dir: Optional[str] = None
threads: Optional[int] = None threads: Optional[int] = None
...@@ -27,7 +29,7 @@ class FastEmbedEncoder(BaseModel): ...@@ -27,7 +29,7 @@ class FastEmbedEncoder(BaseModel):
) )
embedding_args = { embedding_args = {
"model_name": self.model_name, "model_name": self.name,
"max_length": self.max_length, "max_length": self.max_length,
"cache_dir": self.cache_dir, "cache_dir": self.cache_dir,
"threads": self.threads, "threads": self.threads,
......
...@@ -36,11 +36,11 @@ class Encoder: ...@@ -36,11 +36,11 @@ class Encoder:
if self.type == EncoderType.HUGGINGFACE: if self.type == EncoderType.HUGGINGFACE:
raise NotImplementedError raise NotImplementedError
elif self.type == EncoderType.FASTEMBED: elif self.type == EncoderType.FASTEMBED:
self.model = FastEmbedEncoder(name) self.model = FastEmbedEncoder(name=name)
elif self.type == EncoderType.OPENAI: elif self.type == EncoderType.OPENAI:
self.model = OpenAIEncoder(name) self.model = OpenAIEncoder(name=name)
elif self.type == EncoderType.COHERE: elif self.type == EncoderType.COHERE:
self.model = CohereEncoder(name) self.model = CohereEncoder(name=name)
else: else:
raise ValueError raise ValueError
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment