Drop llm.EmbeddingModel.embedding_size property, closes #208

simonw · Sep 2, 2023 · e49c816 · e49c816
1 parent 3d56d6c
commit e49c816
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 5 deletions.
diff --git a/docs/embeddings/writing-plugins.md b/docs/embeddings/writing-plugins.md
@@ -22,14 +22,13 @@ from sentence_transformers import SentenceTransformer
 @llm.hookimpl
 def register_embedding_models(register):
     model_id = "sentence-transformers/all-MiniLM-L6-v2"
-    register(SentenceTransformerModel(model_id, model_id, 384), aliases=("all-MiniLM-L6-v2",))
+    register(SentenceTransformerModel(model_id, model_id), aliases=("all-MiniLM-L6-v2",))
 
 
 class SentenceTransformerModel(llm.EmbeddingModel):
-    def __init__(self, model_id, model_name, embedding_size):
+    def __init__(self, model_id, model_name):
         self.model_id = model_id
         self.model_name = model_name
-        self.embedding_size = embedding_size
         self._model = None
 
     def embed_batch(self, texts):

diff --git a/llm/default_plugins/openai_models.py b/llm/default_plugins/openai_models.py
@@ -63,7 +63,6 @@ def register_embedding_models(register):
 
 class Ada002(EmbeddingModel):
     model_id = "ada-002"
-    embedding_size = 1536
     needs_key = "openai"
     key_env_var = "OPENAI_API_KEY"
     batch_size = 100  # Maybe this should be 2048

diff --git a/llm/models.py b/llm/models.py
@@ -288,7 +288,6 @@ def __repr__(self):
 
 class EmbeddingModel(ABC, _get_key_mixin):
     model_id: str
-    embedding_size: int
     key: Optional[str] = None
     needs_key: Optional[str] = None
     key_env_var: Optional[str] = None