118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
from pymilvus import CollectionSchema, FieldSchema, DataType,client
|
|
from pymilvus import Collection, db, connections,utility,Function,FunctionType
|
|
from pymilvus import (
|
|
MilvusClient, DataType
|
|
)
|
|
|
|
def new_create_coll():
|
|
client = MilvusClient(
|
|
uri="http://172.18.107.78:19530",
|
|
token="root:Milvus"
|
|
)
|
|
client.use_database("JXTest")
|
|
|
|
schema = MilvusClient.create_schema(
|
|
auto_id=True,
|
|
enable_dynamic_field=True,
|
|
)
|
|
|
|
analyzer_params_custom = {
|
|
"tokenizer": "jieba",
|
|
"type": "chinese",
|
|
"filter": ["cnalphanumonly"]
|
|
}
|
|
|
|
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
|
|
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1000,enable_analyzer=True,analyzer_params=analyzer_params_custom)
|
|
schema.add_field(field_name="sparse", datatype=DataType.SPARSE_FLOAT_VECTOR)
|
|
schema.add_field(field_name="dense", datatype=DataType.FLOAT_VECTOR, dim=5)
|
|
|
|
bm25_function = Function(
|
|
name="text_bm25_emb", # Function name
|
|
input_field_names=["text"], # Name of the VARCHAR field containing raw text data
|
|
output_field_names=["sparse"], # Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
|
|
function_type=FunctionType.BM25,
|
|
)
|
|
|
|
schema.add_function(bm25_function)
|
|
|
|
index_params = client.prepare_index_params()
|
|
|
|
# Add indexes
|
|
index_params.add_index(
|
|
field_name="dense",
|
|
index_name="dense_index",
|
|
index_type="AUTOINDEX",#"IVF_FLAT"
|
|
metric_type="COSINE", #L2
|
|
params={"nlist": 128},
|
|
)
|
|
|
|
index_params.add_index(
|
|
field_name="sparse",
|
|
index_name="sparse_index",
|
|
index_type="SPARSE_INVERTED_INDEX", # Index type for sparse vectors
|
|
metric_type="BM25", # Set to `BM25` when using function to generate sparse vectors
|
|
params={"inverted_index_algo": "DAAT_MAXSCORE"}, # The ratio of small vector values to be dropped during indexing
|
|
)
|
|
|
|
client.create_collection(
|
|
collection_name="hybrid_search_collection",
|
|
schema=schema,
|
|
index_params=index_params
|
|
)
|
|
return "Sucessful"
|
|
|
|
|
|
|
|
def create_collections(collections_name,database_name : str ="JXTest"):
|
|
db.using_database(database_name)
|
|
|
|
id = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True)
|
|
text = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=4096,enable_analyzer=True)
|
|
header_1 = FieldSchema(name="header_1", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
|
header_2 = FieldSchema(name="header_2", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
|
header_3 = FieldSchema(name="header_3", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
|
vector = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024,)
|
|
sparse = FieldSchema(name="sparse_bm25", dtype=DataType.SPARSE_FLOAT_VECTOR)
|
|
schema = CollectionSchema(
|
|
fields=[id, header_1, header_2, header_3, text, vector,sparse],
|
|
auto_id=True,
|
|
enable_dynamic_field=True
|
|
)
|
|
index_params = {
|
|
"field_name": "vector",
|
|
"index_name": "vector_index",
|
|
"metric_type": "L2",
|
|
"index_type": "IVF_FLAT", #FLAT
|
|
"params": {"nlist": 128},
|
|
"field_name":"sparse_bm25",
|
|
"index_name":"sparse_bm25_index",
|
|
"index_type":"SPARSE_WAND",
|
|
"metric_type":"BM25",
|
|
"params":{
|
|
"inverted_index_algo": "DAAT_MAXSCORE", # Algorithm for building and querying the index. Valid values: DAAT_MAXSCORE, DAAT_WAND, TAAT_NAIVE.
|
|
"bm25_k1": 1.2,
|
|
"bm25_b": 0.75
|
|
}
|
|
|
|
}
|
|
bm25_function = Function(
|
|
name="bm25",
|
|
function_type=FunctionType.BM25,
|
|
input_field_names=["text"],
|
|
output_field_names="sparse_bm25",
|
|
)
|
|
schema.add_function(bm25_function)
|
|
|
|
collection = Collection(name=collections_name, schema=schema, using='default', shards_num=2)
|
|
collection.create_index(field_name="vector", index_params=index_params)
|
|
|
|
utility.index_building_progress("Hybrid")
|
|
return "collections created"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# conn = connections.connect(host="172.18.107.78",port=19530,user="user",password="Milvus")
|
|
# res = create_collections("HybirdTest","JXTest")
|
|
# print(res)
|
|
res = new_create_coll() |