Skip to content
This repository has been archived by the owner on Feb 22, 2020. It is now read-only.

Commit

Permalink
fix(service): add doc type to req generator
Browse files Browse the repository at this point in the history
  • Loading branch information
Larryjianfeng committed Jul 26, 2019
1 parent 80e234e commit d834f57
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions gnes/proto/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

class RequestGenerator:
@staticmethod
def index(data: List[bytes], batch_size: int = 0, start_id: int = 0, doc_type: str = 'TEXT', *args, **kwargs):
def index(data: List[bytes], batch_size: int = 0, start_id: int = 0, doc_type: int = 1, *args, **kwargs):

for pi in batch_iterator(data, batch_size):
req = gnes_pb2.Request()
Expand All @@ -42,7 +42,7 @@ def index(data: List[bytes], batch_size: int = 0, start_id: int = 0, doc_type: s
start_id += 1

@staticmethod
def train(data: List[bytes], batch_size: int = 0, start_id: int = 0, doc_type: str = 'TEXT', *args, **kwargs):
def train(data: List[bytes], batch_size: int = 0, start_id: int = 0, doc_type: int = 1, *args, **kwargs):
for pi in batch_iterator(data, batch_size):
req = gnes_pb2.Request()
req.request_id = str(start_id)
Expand All @@ -59,7 +59,7 @@ def train(data: List[bytes], batch_size: int = 0, start_id: int = 0, doc_type: s
start_id += 1

@staticmethod
def query(query: bytes, top_k: int, start_id: int = 0, doc_type: str = 'TEXT', *args, **kwargs):
def query(query: bytes, top_k: int, start_id: int = 0, doc_type: int = 1, *args, **kwargs):
if top_k <= 0:
raise ValueError('"top_k: %d" is not a valid number' % top_k)

Expand Down

0 comments on commit d834f57

Please sign in to comment.