text_to_vector/utils/milvus_hlper.py

# -*- coding: utf-8 -*-
# @Time    : 2023/10/11 
# @Author  : lkj
# @description : milvus
from typing import Optional, List

import numpy
from numpy import dot
from numpy.linalg import norm
from pymilvus import (connections, FieldSchema,Collection,CollectionSchema,DataType,utility)
from config import redis_config
from utils.redis_helper import RedisString



# 使用默认数据库 ‘default’,也可以自己建数据库


class Milvus(object):
    def __init__(self,table,**kwargs):

        connections.connect(**kwargs)
        self.col = Collection(table)
        self.search_params = {
            "metric_type": "L2",
            "ignore_growing": False,
            "params": {"nprobe": 100},
        }
        self.r = RedisString(redis_config)

    def search(self,vec,fileds=None,expr=None):
        res = self.col.search([vec], 'embeddings', self.search_params, 7,
                              output_fields=fileds,
                              expr=expr)
        return res

    def load(self):
        self.col.load()

    def release(self):
        self.col.release()

    def delete(self,expr):
        self.col.delete(expr=expr)

    def query(self,q):
        res = self.col.query(expr=q)
        return res

    def insert(self,data):
        self.col.insert(data=data)

    @staticmethod
    def sim(a: list, b: list):
        """
        余弦计算两个向量相似度
        :param a:
        :param b:
        :return:
        """
        s = dot(a, b) / (norm(a) * norm(b))
        return round(s, 4)

    def get_name(self,code):
        """"
        基于redis查询code对应name
        """
        while code[-1] == '0' and code[-2] == '0':
            code = code[:-1]
            code = code[:-1]
        name = self.r.string_get('jycode_' + code)
        return name

    def get_root_zc(self, re_code):
        """
        根据code查询对应root
        """
        split = 0
        root = ''
        level = (len(re_code) - 1) / 2
        for i in range(int(level)):
            if split == 0:
                c = re_code
            else:
                c = re_code[:-split]
            split += 2
            name_code = self.get_name(c)
            root = name_code + '/' + root
        return root

    def search_good(self,vec,num=7,base=None):
        try:
            res = self.col.search([vec],
                                  'embeddings', self.search_params, num,
                                  output_fields=['code','class_name','embeddings','explain','root'],) #expr=f'baseclass=="{base}"'
            result_list = []
            for hit in res:
                # print(hit)
                for row in hit:
                    row = row.to_dict()
                    code = row.get('entity',{}).get('code','')
                    while code[-1] == '0' and code[-2] == '0':
                        code = code[:-1]
                        code = code[:-1]
                    explain = row.get('entity', {}).get('explain', '')
                    root = self.get_root_zc(code)
                    if not root:
                        root = row.get('entity', {}).get('root', '')
                    vec_cls = row.get('entity', {}).get('embeddings', [])
                    sim_res = self.sim(vec, vec_cls)
                    cls_name = row.get('entity', {}).get('class_name', '')
                    result_list.append((cls_name, code, sim_res, root, explain))
            result_list = sorted(result_list, key=lambda x: x[2], reverse=True)
            return result_list
        except Exception as e:
            print('关系库错误:', e)
            return []

    def search_industry(self, vec, output_fields: list, topk=7,
                        industry_list: Optional[List[str]] = None,):
        """
        查询统计局分类函数
        vec ： 标的物转成的向量
        industry_list：行业范围 -->list

        """
        try:
            public = True
            if industry_list:
                # expr = f'industry in {industry_list}'
                public = False
            else:
                expr = None
            res = self.col.search([vec], 'embeddings', self.search_params, topk,
                                  output_fields=output_fields,
                                  expr=None)
            result_list = []
            for hit in res:
                for row in hit:
                    row = row.to_dict()
                    code = row.get('entity', {}).get('code', '')
                    cls_name = row.get('entity', {}).get('class_name', '')
                    if not public:
                        code = row.get('entity', {}).get('private_code', '')
                        if code == 'null':
                            code = row.get('entity', {}).get('code', '')
                        cls_name = self.get_name(code)
                    while code[-1] == '0' and code[-2] == '0':
                        code = code[:-1]
                        code = code[:-1]
                    explain = row.get('entity', {}).get('explain', '')
                    root = self.get_root_zc(code)
                    if not root:
                        root = row.get('entity', {}).get('root', '')
                    vec_cls = row.get('entity', {}).get('embeddings', [])
                    sim_res = self.sim(vec, vec_cls)
                    result_list.append((cls_name, code, sim_res, root, explain))
            result_list = sorted(result_list, key=lambda x: x[2], reverse=True)
            return result_list
        except Exception as e:
            print('统计局分类错误:', e)
            return []

    def search_china(self, vec,base=None,topk=7):
        """
        查询统计局分类函数
        vec ： 标的物转成的向量
        base：标的物分类

        """
        try:

            res = self.col.search([vec], 'embeddings', self.search_params, topk,
                                  output_fields=['code','class_name','embeddings','explain','root'],
                                  )#expr=f'baseclass=="{base}"'
            result_list = []
            for hit in res:
                # print(hit)
                for row in hit:

                    row = row.to_dict()
                    code = row.get('entity',{}).get('code','')
                    while code[-1] == '0' and code[-2] == '0':
                        code = code[:-1]
                        code = code[:-1]
                    ids = row['id']
                    explain = row.get('entity', {}).get('explain', '')
                    root = self.get_root_zc(code)
                    if not root:
                        root = row.get('entity', {}).get('root', '')
                    vec_cls = row.get('entity', {}).get('embeddings', [])
                    sim_res = self.sim(vec, vec_cls)
                    cls_name = row.get('entity', {}).get('class_name', '')
                    result_list.append((cls_name, code, sim_res, root, explain))
            result_list = sorted(result_list, key=lambda x: x[2], reverse=True)
            return result_list
        except Exception as e:
            print('统计局分类错误:', e)
            return []

    def update(self, data):
        self.col.insert(data)

def create():
    onn = connections.connect(db_name="classify", host="192.168.3.109", port=19530)
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="class_name", dtype=DataType.VARCHAR, max_length=100),
        FieldSchema(name='code', dtype=DataType.VARCHAR, max_length=16),
        FieldSchema(name="p_name", dtype=DataType.VARCHAR, max_length=100),
        FieldSchema(name='p_code', dtype=DataType.VARCHAR, max_length=16),
        # FieldSchema(name='baseclass', dtype=DataType.VARCHAR, max_length=10),
        FieldSchema(name='explain', dtype=DataType.VARCHAR, max_length=500),
        # FieldSchema(name='root', dtype=DataType.VARCHAR, max_length=300),
        FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=768),
    ]

    schema = CollectionSchema(fields, description='场景分类')
    milvus_conn = Collection('scene', schema=schema)
    index = {"index_type": "IVF_FLAT", "metric_type": "L2",  "params":{"nlist":128}}
    milvus_conn.create_index("embeddings", index)

if __name__ == '__main__':
    # create()
    # pass
    from request_fun import text_to_vector
    from config import milvus_config
    text = '施工人员派遣'
    vector = text_to_vector(text)
    col = Milvus('jianyu_code',**milvus_config)
    print(col.search_industry(vector, ['code', 'class_name', 'embeddings', 'explain', 'root', 'private_code'],
                              industry_list=['物业']))