From 5006b49ba8952ca44e1e9698f3926c507f48aae8 Mon Sep 17 00:00:00 2001
From: James Briggs <35938317+jamescalam@users.noreply.github.com>
Date: Sat, 6 Jan 2024 16:53:58 +0100
Subject: [PATCH] lint fixes
---
semantic_router/encoders/fastembed.py | 10 ++++++----
semantic_router/schema.py | 6 +++---
2 files changed, 9 insertions(+), 7 deletions(-)
diff --git a/semantic_router/encoders/fastembed.py b/semantic_router/encoders/fastembed.py
index c730d4c6..fb845ce7 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 826cb686..64480355 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
--
GitLab