Use type specific model for other search types too. Expose them via SearchModels

- Wrap Image, Music, Ledger search into the type of SearchModel they use
  Similar to what was done for notes model by wrapping it's config
  into an AsymmetricSearchModel.

- Use the uber wrapper class to expose all type specific search models
This commit is contained in:
Debanjum Singh Solanky
2021-09-29 21:09:42 -07:00
parent 352d2930ee
commit f4dd9cd117
2 changed files with 32 additions and 37 deletions

View File

@@ -27,6 +27,26 @@ class AsymmetricSearchModel():
self.top_k = top_k
class LedgerSearchModel():
def __init__(self, transactions, transaction_embeddings, symmetric_encoder, symmetric_cross_encoder, top_k):
self.transactions = transactions
self.transaction_embeddings = transaction_embeddings
self.symmetric_encoder = symmetric_encoder
self.symmetric_cross_encoder = symmetric_cross_encoder
self.top_k = top_k
class ImageSearchModel():
def __init__(self, image_names, image_embeddings, image_metadata_embeddings, image_encoder):
self.image_names = image_names
self.image_embeddings = image_embeddings
self.image_metadata_embeddings = image_metadata_embeddings
self.image_encoder = image_encoder
@dataclass
class SearchModels():
notes_search: AsymmetricSearchModel = None
ledger_search: LedgerSearchModel = None
music_search: AsymmetricSearchModel = None
image_search: ImageSearchModel = None