From 8eba19dc59c82851047eec1d5ebf557ed5d080a2 Mon Sep 17 00:00:00 2001 From: DylanAlloy Date: Fri, 15 Mar 2024 08:35:13 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=A7=20fix(wip):=20good=20milvus=20chan?= =?UTF-8?q?ges?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- magnet/ize/memory.py | 8 ++++---- magnet/utils/index/milvus.py | 12 ++++++------ 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/magnet/ize/memory.py b/magnet/ize/memory.py index 6e097b8..314ab29 100644 --- a/magnet/ize/memory.py +++ b/magnet/ize/memory.py @@ -31,12 +31,12 @@ async def on(self, create: bool = False, initialize: bool = False): self._model = SentenceTransformer(self.config.index.model, device=Utils().check_cuda()) _f('info', f'loading into {self._model.device}') self.db = MilvusDB(self.config) - self.db.on() + await self.db.on() if create: - self.db.create(overwrite=True) - self.db.load() + await self.db.create(overwrite=True) + await self.db.load() if initialize: - self.db.initialize() + self.db.initialize() async def index(self, payload, msg, field: Charge = None, v: bool = False, instruction: str = "Represent this sentence for searching relevant passages: "): if not msg or not payload: diff --git a/magnet/utils/index/milvus.py b/magnet/utils/index/milvus.py index e24f608..79ca60a 100644 --- a/magnet/utils/index/milvus.py +++ b/magnet/utils/index/milvus.py @@ -12,7 +12,7 @@ def __init__(self, config: MagnetConfig): FieldSchema(name='text', dtype=DataType.VARCHAR, max_length=65535), FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, dim=self.config.index.dimension) ] - def on(self): + async def on(self): try: _f('wait', f'connecting to {self.config.index.milvus_uri}') self.connection = connections.connect( @@ -34,17 +34,17 @@ async def off(self): except Exception as e: return _f('fatal', e) - def create(self, overwrite=False): - if utility.has_collection(self.config.index.name, using=self.config.session) and overwrite: - utility.drop_collection(self.config.index.name, using=self.config.session) + async def create(self, overwrite=False): try: + if utility.has_collection(self.config.index.name, using=self.config.session) and overwrite: + utility.drop_collection(self.config.index.name, using=self.config.session) self.collection = Collection(name=self.config.index.name, schema=self.schema, using=self.config.session) self.collection.create_index(field_name="embedding", index_params=self.config.index.options) _f('success', f"{self.config.index.name} created") except Exception as e: _f('fatal', e) - def load(self): + async def load(self): _f('wait', f'loading {self.config.index.name} into memory, may take time') self.collection = Collection(name=self.config.index.name, schema=self.schema, using=self.config.session) self.collection.load() @@ -71,7 +71,7 @@ def initialize(self, user: str = 'magnet', password: str = '33011033'): except Exception as e: _f('fatal', e) - def delete_index(self): + async def delete_index(self): if utility.has_collection(self.config.index.name, using=self.config.session): utility.drop_collection(self.config.index.name, using=self.config.session) _f('warn', f"Index for {self.config.index.name} deleted")