def extract_features(filepath, model): try: cache = Cache('./tmp') feats = [] names = [] total = len(open(filepath, 'rU').readlines()) cache['total'] = total current = 0 vec = 0 with open(filepath, 'r') as f: for line in f: current += 1 cache['current'] = current line = line.strip() line = line.split() line = line[0] try: vec = smiles_to_vector(line) feats.append(vec) names.append(line.encode()) print( "extracted feature from smi No. %d , %d molecular in total" % (current, total)) except Exception as e: print( "failed to extract feature from smi No. %d , %d molecular in total" % (current, total), e) return feats, names except Exception as e: LOGGER.error(" Error with extracting feature from image {}".format(e)) sys.exit(1)
def has_collection(self, collection_name): try: status = self.client.has_collection(collection_name)[1] return status except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1)
def audio_path(audio_path): # Get the audio file try: LOGGER.info(("Successfully load audio: {}".format(audio_path))) return FileResponse(audio_path) except Exception as e: LOGGER.error("upload audio error: {}".format(e)) return {'status': False, 'msg': e}, 400
def delete_table(self, table_name): sql = "drop table if exists " + table_name + ";" try: self.cursor.execute(sql) LOGGER.debug("MYSQL delete table:{}".format(table_name)) except Exception as e: LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql)) sys.exit(1)
def get_progress(): # Get the progress of dealing with data try: cache = Cache('./tmp') return "current: {}, total: {}".format(cache['current'], cache['total']) except Exception as e: LOGGER.error("upload data error: {}".format(e)) return {'status': False, 'msg': e}, 400
def mols_img(mols_path): # Get the molecular image file try: LOGGER.info(("Successfully load molecular image: {}".format(mols_path))) return FileResponse(UPLOAD_PATH + '/' + mols_path + '.png') except Exception as e: LOGGER.error("upload image error: {}".format(e)) return {'status': False, 'msg': e}, 400
async def count_audio(table_name: str = None): # Returns the total number of vectors in the system try: num = do_count(table_name, MILVUS_CLI) LOGGER.info("Successfully count the number of data!") return num except Exception as e: LOGGER.error(e) return {'status': False, 'msg': e}, 400
async def load_data(item: Item): # Insert all the data under the file path to Milvus/MySQL try: total_num = do_load(item.Table, item.File, MODEL, MILVUS_CLI, MYSQL_CLI) LOGGER.info("Successfully loaded data, total count: {}".format(total_num)) return {'status': True, 'msg': "Successfully loaded data!"} except Exception as e: LOGGER.error(e) return {'status': False, 'msg': e}, 400
async def drop_tables(table_name: str = None): # Delete the collection of Milvus and MySQL try: status = do_drop(table_name, MILVUS_CLI, MYSQL_CLI) LOGGER.info("Successfully drop tables in Milvus and MySQL!") return status except Exception as e: LOGGER.error(e) return {'status': False, 'msg': e}, 400
def __init__(self): try: self.client = Milvus(host=MILVUS_HOST, port=MILVUS_PORT) LOGGER.debug( "Successfully connect to Milvus with IP:{} and PORT:{}".format( MILVUS_HOST, MILVUS_PORT)) except Exception as e: LOGGER.error("Failed to connect Milvus: {}".format(e)) sys.exit(1)
def delete_collection(self, collection_name): try: self.set_collection(collection_name) self.collection.drop() LOGGER.debug("Successfully drop collection!") return "ok" except Exception as e: LOGGER.error("Failed to drop collection: {}".format(e)) sys.exit(1)
def create_mysql_table(self, table_name): sql = "create table if not exists " + table_name + "(milvus_id TEXT, data_path TEXT);" try: self.cursor.execute(sql) LOGGER.debug("MYSQL create table: {} with sql: {}".format( table_name, sql)) except Exception as e: LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql)) sys.exit(1)
def count_table(self, table_name): sql = "select count(milvus_id) from " + table_name + ";" try: self.cursor.execute(sql) results = self.cursor.fetchall() LOGGER.debug("MYSQL count table:{}".format(table_name)) return results[0][0] except Exception as e: LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql)) sys.exit(1)
def delete_all_data(self, table_name): sql = 'delete from ' + table_name + ';' try: self.cursor.execute(sql) self.conn.commit() LOGGER.debug( "MYSQL delete all data in table:{}".format(table_name)) except Exception as e: LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql)) sys.exit(1)
def do_drop(table_name, milvus_cli, mysql_cli): if not table_name: table_name = DEFAULT_TABLE try: status = milvus_cli.delete_collection(table_name) mysql_cli.delete_table(table_name) return status except Exception as e: LOGGER.error(" Error with drop table: {}".format(e)) sys.exit(1)
def load_data_to_mysql(self, table_name, data): sql = "insert into " + table_name + " (milvus_id,data_path) values (%s,%s);" try: self.cursor.executemany(sql, data) self.conn.commit() LOGGER.debug("MYSQL loads data to table: {} successfully".format( table_name)) except Exception as e: LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql)) sys.exit(1)
def set_collection(self, collection_name): try: if self.has_collection(collection_name): self.collection = Collection(name=collection_name) else: raise Exception( "There has no collection named:{}".format(collection_name)) except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1)
def count(self, collection_name): try: self.set_collection(collection_name) num = self.collection.num_entities LOGGER.debug( "Successfully get the num:{} of the collection:{}".format( num, collection_name)) return num except Exception as e: LOGGER.error("Failed to count vectors in Milvus: {}".format(e)) sys.exit(1)
def do_count(table_name, milvus_cli): if not table_name: table_name = DEFAULT_TABLE try: if not milvus_cli.has_collection(table_name): return None num = milvus_cli.count(table_name) return num except Exception as e: LOGGER.error(" Error attempting to count table {}".format(e)) sys.exit(1)
def search_by_milvus_ids(self, ids, table_name): str_ids = str(ids).replace('[', '').replace(']', '') sql = "select data_path from " + table_name + " where milvus_id in (" + str_ids + ") order by field (milvus_id," + str_ids + ");" try: self.cursor.execute(sql) results = self.cursor.fetchall() results = [res[0] for res in results] LOGGER.debug("MYSQL search by milvus id.") return results except Exception as e: LOGGER.error("MYSQL ERROR: {} with sql: {}".format(e, sql)) sys.exit(1)
def do_drop(table_name, milvus_cli, mysql_cli): if not table_name: table_name = DEFAULT_TABLE try: if not milvus_cli.has_collection(table_name): return "collection is not exist" status = milvus_cli.delete_collection(table_name) mysql_cli.delete_table(table_name) return status except Exception as e: LOGGER.error(" Error attempting to drop table: {}".format(e)) sys.exit(1)
def do_drop(table_name, milvus_cli, mysql_cli): if not table_name: table_name = DEFAULT_TABLE try: mysql_cli.delete_table(table_name) if not milvus_cli.has_collection(table_name): raise Exception("When drop table, there has no table named " + table_name) status = milvus_cli.delete_collection(table_name) return status except Exception as e: LOGGER.error(" Error with drop table: {}".format(e)) sys.exit(1)
async def search_data(Table: str = None, Mol: str = None): # Search the upload image in Milvus/MySQL try: # Save the upload data to server. ids, paths, distances = do_search(Table, Mol, MODEL, MILVUS_CLI, MYSQL_CLI) res = dict(zip(paths, zip(ids, distances))) res = sorted(res.items(), key=lambda item: item[1]) LOGGER.info("Successfully searched similar data!") return res except Exception as e: LOGGER.error(e) return {'status': False, 'msg': e}, 400
def delete_collection(self, collection_name): try: status = self.client.drop_collection( collection_name=collection_name) if not status.code: LOGGER.debug( "Successfully drop collection: {}".format(collection_name)) return status else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to drop collection: {}".format(e)) sys.exit(1)
def insert(self, collection_name, vectors): try: self.create_collection(collection_name) data = [vectors] mr = self.collection.insert(data) ids = mr.primary_keys self.collection.load() LOGGER.debug( "Insert vectors to Milvus in collection: {} with {} rows". format(collection_name, len(vectors))) return ids except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1)
def do_search(table_name, molecular_name, model, milvus_client, mysql_cli): if not table_name: table_name = DEFAULT_TABLE try: feat = smiles_to_vector(molecular_name) print(feat) vectors = milvus_client.search_vectors(table_name, [feat], TOP_K) vids = [str(x.id) for x in vectors[0]] smiles = mysql_cli.search_by_milvus_ids(vids, table_name) distances = [x.distance for x in vectors[0]] return vids, smiles, distances except Exception as e: LOGGER.error(" Error with search : {}".format(e)) sys.exit(1)
def count(self, collection_name): try: status, num = self.client.count_entities( collection_name=collection_name) if not status.code: LOGGER.debug( "Successfully get the num:{} of the collection:{}".format( num, collection_name)) return num else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to count vectors in Milvus: {}".format(e)) sys.exit(1)
def do_count(table_name, milvus_cli, mysql_cli): if not table_name: table_name = DEFAULT_TABLE try: if not milvus_cli.has_collection(table_name): return None milvus_num = milvus_cli.count(table_name) mysql_num = mysql_cli.count_table(table_name) LOGGER.debug("The num of Milvus: {} and Mysql: {}".format( milvus_num, mysql_num)) assert milvus_num == mysql_num return milvus_num except Exception as e: LOGGER.error(" Error with count table {}".format(e)) sys.exit(1)
def create_index(self, collection_name): try: index_param = {'nlist': 16384} status = self.client.create_index(collection_name, IndexType.IVF_FLAT, index_param) if not status.code: LOGGER.debug( "Successfully create index in collection:{} with param:{}". format(collection_name, index_param)) return status else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to create index: {}".format(e)) sys.exit(1)
def insert(self, collection_name, vectors): try: self.create_colllection(collection_name) status, ids = self.client.insert(collection_name=collection_name, records=vectors) if not status.code: LOGGER.debug( "Insert vectors to Milvus in collection: {} with {} rows". format(collection_name, len(vectors))) return ids else: raise Exception(status.message) except Exception as e: LOGGER.error("Failed to load data to Milvus: {}".format(e)) sys.exit(1)