Skip to content

Commit

Permalink
Add default configuration for Embeddings, closes #393
Browse files Browse the repository at this point in the history
  • Loading branch information
davidmezzetti committed Dec 3, 2022
1 parent c90d781 commit 3707daa
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 0 deletions.
14 changes: 14 additions & 0 deletions src/python/txtai/embeddings/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,10 @@ def index(self, documents, reindex=False):
reindex: if this is a reindex operation in which case database creation is skipped, defaults to False
"""

# Set configuration to default configuration, if empty
if not self.config:
self.configure(self.defaults())

# Create document database, if necessary
if not reindex:
self.database = self.createdatabase()
Expand Down Expand Up @@ -571,6 +575,16 @@ def configure(self, config):
# Query model
self.query = self.loadquery() if self.config else None

def defaults(self):
"""
Builds a default configuration.
Returns:
default configuration
"""

return {"path": "sentence-transformers/all-MiniLM-L6-v2"}

def loadvectors(self):
"""
Loads a vector model set in config.
Expand Down
11 changes: 11 additions & 0 deletions test/python/testembeddings.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,17 @@ def tearDownClass(cls):
if cls.embeddings:
cls.embeddings.close()

def testDefaults(self):
"""
Test default configuration
"""

# Run index with no config which will fall back to default configuration
embeddings = Embeddings()
embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])

self.assertEqual(embeddings.count(), 6)

def testDelete(self):
"""
Test delete
Expand Down

0 comments on commit 3707daa

Please sign in to comment.