Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update elastic connection #62

Merged
merged 1 commit into from
Aug 11, 2023
Merged

Update elastic connection #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,10 +85,12 @@
# Scalar db configs
SCALARDB_CONFIG = {
'connection_args': {
'hosts': os.getenv('ES_HOSTS', 'https://localhost:9200'),
# 'hosts': os.getenv('ES_HOSTS', 'https://localhost:9200'),
'cloud_id': os.getenv('ES_CLOUD_ID'),
'ca_certs': os.getenv('ES_CA_CERTS', None),
'basic_auth': (os.getenv('ES_USER', 'user_name'), os.getenv('ES_PASSWORD', 'es_password'))
},
'top_k': 3
}

# Memory db configs
Expand Down
2 changes: 1 addition & 1 deletion requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,4 +10,4 @@ fastapi
uvicorn
towhee>=1.1.0
pymilvus
elasticsearch
elasticsearch>=8.0.0
39 changes: 13 additions & 26 deletions src_towhee/pipelines/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,18 +66,9 @@ def __init__(self,
if self.use_scalar:
from elasticsearch import Elasticsearch # pylint: disable=C0415

self.es_uri = scalardb_config['connection_args']['hosts']
self.es_host = self.es_uri.split('https://')[1].split(':')[0]
self.es_port = self.es_uri.split('https://')[1].split(':')[1]
self.es_ca_certs = scalardb_config['connection_args']['ca_certs']
self.es_basic_auth = scalardb_config['connection_args']['basic_auth']
self.es_user = self.es_basic_auth[0]
self.es_password = self.es_basic_auth[1]
self.es_client = Elasticsearch(
self.es_uri,
ca_certs=self.es_ca_certs,
basic_auth=self.es_basic_auth
)
self.es_connection_kwargs = scalardb_config['connection_args']
self.es_top_k = scalardb_config['top_k']
self.es_client = Elasticsearch(**self.es_connection_kwargs)

@property
def search_pipeline(self):
Expand Down Expand Up @@ -117,11 +108,8 @@ def search_config(self):
# Configure scalar store (ES)
if self.use_scalar:
search_config.es_enable = True
search_config.es_host = self.es_host
search_config.es_port = self.es_port
search_config.es_user = self.es_user
search_config.es_password = self.es_password
search_config.es_ca_certs = self.es_ca_certs
search_config.es_connection_kwargs = self.es_connection_kwargs
search_config.es_top_k = self.es_top_k
else:
search_config.es_enable = False
return search_config
Expand Down Expand Up @@ -150,11 +138,7 @@ def insert_config(self):
# Configure scalar store (ES)
if self.use_scalar:
insert_config.es_enable = True
insert_config.es_host = self.es_host
insert_config.es_port = self.es_port
insert_config.es_user = self.es_user
insert_config.es_password = self.es_password
insert_config.es_ca_certs = self.es_ca_certs
insert_config.es_connection_kwargs = self.es_connection_kwargs
else:
insert_config.es_enable = False
return insert_config
Expand All @@ -168,11 +152,14 @@ def create(self, project: str):
FieldSchema(name='text_id', dtype=DataType.VARCHAR,
description='text', max_length=500),
FieldSchema(name='text', dtype=DataType.VARCHAR,
description='text', max_length=1000),
FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR,
description='embedding vectors', dim=self.textencoder_config['dim'])
description='text', max_length=1000)
]
schema = CollectionSchema(fields=fields, description='osschat', enable_dynamic_field=True)
if INSERT_MODE == 'generate_questions':
fields.append(FieldSchema(name='question', dtype=DataType.VARCHAR,
description='generated question', max_length=500))
fields.append(FieldSchema(name='embedding', dtype=DataType.FLOAT_VECTOR,
description='embedding vectors', dim=self.textencoder_config['dim']))
schema = CollectionSchema(fields=fields, description='osschat', enable_dynamic_field=False)
collection = Collection(name=project, schema=schema)

index_params = self.milvus_index_params
Expand Down
11 changes: 3 additions & 8 deletions src_towhee/pipelines/insert/generate_questions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,22 +44,17 @@ def custom_pipeline(config):
if config.embedding_normalize:
p = p.map('embedding', 'embedding', ops.towhee.np_normalize())

p = p.map(('project', 'chunk', 'gen_question', 'embedding'), 'milvus_res',
p = p.map(('project', 'doc', 'chunk', 'gen_question', 'embedding'), 'milvus_res',
ops.ann_insert.osschat_milvus(host=config.milvus_host,
port=config.milvus_port,
user=config.milvus_user,
password=config.milvus_password,
))

if config.es_enable:
es_index_op = ops.elasticsearch.osschat_index(host=config.es_host,
port=config.es_port,
user=config.es_user,
password=config.es_password,
ca_certs=config.es_ca_certs,
)
es_index_op = ops.elasticsearch.osschat_index(**config.es_connection_kwargs)
p = (
p.map(('gen_question', 'chunk'), 'es_doc', lambda x, y: {'sentence': x, 'doc': y})
p.map(('gen_question', 'chunk'), 'es_doc', lambda x, y: {'question': x, 'doc': y})
.map(('project', 'es_doc'), 'es_res', es_index_op)
.map(('milvus_res', 'es_res'), 'res', lambda x, y: {'milvus_res': x, 'es_res': y})
)
Expand Down
Loading