通过阿里云向量检索 Milvus 版和通义千问快速构建基于专属知识库的问答系统

pip3 install pymilvus tqdm dashscope

git clone https://github.com/shijiebei2009/CEC-Corpus.git

import osimport timefrom tqdm import tqdmimport dashscopefrom dashscope import TextEmbeddingfrom pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility

def prepareData(path, batch_size=25):    batch_docs = []    for file in os.listdir(path):        with open(path + '/' + file, 'r', encoding='utf-8') as f:            batch_docs.append(f.read())            if len(batch_docs) == batch_size:                yield batch_docs                batch_docs = []
    if batch_docs:        yield batch_docs

def getEmbedding(news):    model = TextEmbedding.call(        model=TextEmbedding.Models.text_embedding_v1,        input=news    )    embeddings = [record['embedding'] for record in model.output['embeddings']]    return embeddings if isinstance(news, list) else embeddings[0]

if __name__ == '__main__':
    current_path = os.path.abspath(os.path.dirname(__file__))   # 当前目录    root_path = os.path.abspath(os.path.join(current_path, '..'))   # 上级目录    data_path = f'{root_path}/CEC-Corpus/raw corpus/allSourceText'  # 数据下载git clone https://github.com/shijiebei2009/CEC-Corpus.git
    # 配置Dashscope API KEY    dashscope.api_key = '<YOUR_DASHSCOPE_API_KEY>'
    # 配置Milvus参数    COLLECTION_NAME = 'CEC_Corpus'    DIMENSION = 1536    MILVUS_HOST = 'c-97a7d8038fb8****.milvus.aliyuncs.com'    MILVUS_PORT = '19530'    USER = 'root'    PASSWORD = '<password>'
    connections.connect(host=MILVUS_HOST, port=MILVUS_PORT, user=USER, password=PASSWORD)
    # Remove collection if it already exists    if utility.has_collection(COLLECTION_NAME):        utility.drop_collection(COLLECTION_NAME)
    # Create collection which includes the id, title, and embedding.    fields = [        FieldSchema(name='id', dtype=DataType.INT64, descrition='Ids', is_primary=True, auto_id=False),        FieldSchema(name='text', dtype=DataType.VARCHAR, description='Text', max_length=4096),        FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, description='Embedding vectors', dim=DIMENSION)    ]    schema = CollectionSchema(fields=fields, description='CEC Corpus Collection')    collection = Collection(name=COLLECTION_NAME, schema=schema)
    # Create an index for the collection.    index_params = {        'index_type': 'IVF_FLAT',        'metric_type': 'L2',        'params': {'nlist': 1024}    }    collection.create_index(field_name="embedding", index_params=index_params)
    id = 0    for news in tqdm(list(prepareData(data_path))):        ids = [id + i for i, _ in enumerate(news)]        id += len(news)
        vectors = getEmbedding(news)        # insert Milvus Collection        for id, vector, doc in zip(ids, vectors, news):            insert_doc = (doc[:498] + '..') if len(doc) > 500 else doc            ins = [[id], [insert_doc], [vector]]  # Insert the title id, the text, and the text embedding vector            collection.insert(ins)            time.sleep(2)

import osimport dashscopefrom dashscope import Generationfrom pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collectionfrom embedding import getEmbedding

def getAnswer(query, context):    prompt = f'''请基于```内的报道内容，回答我的问题。	      ```	      {context}	      ```	      我的问题是：{query}。       '''
    rsp = Generation.call(model='qwen-turbo', prompt=prompt)    return rsp.output.text

def search(text):    # Search parameters for the index    search_params = {        "metric_type": "L2"    }
    results = collection.search(        data=[getEmbedding(text)],  # Embeded search value        anns_field="embedding",  # Search across embeddings        param=search_params,        limit=1,  # Limit to five results per search        output_fields=['text']  # Include title field in result    )
    ret = []    for hit in results[0]:        ret.append(hit.entity.get('text'))    return ret

if __name__ == '__main__':
    current_path = os.path.abspath(os.path.dirname(__file__))   # 当前目录    root_path = os.path.abspath(os.path.join(current_path, '..'))   # 上级目录    data_path = f'{root_path}/CEC-Corpus/raw corpus/allSourceText'
    # 配置Dashscope API KEY    dashscope.api_key = '<YOUR_DASHSCOPE_API_KEY>'
    # 配置Milvus参数    COLLECTION_NAME = 'CEC_Corpus'    DIMENSION = 1536    MILVUS_HOST = 'c-97a7d8038fb8****.milvus.aliyuncs.com'    MILVUS_PORT = '19530'    USER = 'root'    PASSWORD = '<password>'
    connections.connect(host=MILVUS_HOST, port=MILVUS_PORT, user=USER, password=PASSWORD)
    fields = [        FieldSchema(name='id', dtype=DataType.INT64, descrition='Ids', is_primary=True, auto_id=False),        FieldSchema(name='text', dtype=DataType.VARCHAR, description='Text', max_length=4096),        FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR, description='Embedding vectors', dim=DIMENSION)    ]    schema = CollectionSchema(fields=fields, description='CEC Corpus Collection')    collection = Collection(name=COLLECTION_NAME, schema=schema)
    # Load the collection into memory for searching    collection.load()
    question = '北京中央电视台工地发生大火，发生在哪里？出动了多少辆消防车？人员伤亡情况如何？'    context = search(question)    answer = getAnswer(question, context)    print(answer)