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

lint fixes

parent c9715a34
No related branches found
No related tags found
No related merge requests found
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,
......
......@@ -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
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment