Fix classify and classify_batch for Python client (#608)

Author: tgaddair

clients/python/lorax/client.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict, Optional, List, AsyncIterator, Iterator, Union
 
 from lorax.types import (
+    BatchRequest,
     StreamResponse,
     Response,
     Request,
@@ -69,6 +70,7 @@ def __init__(
         self.base_url = base_url
         self.embed_endpoint = f"{base_url}/embed"
         self.classify_endpoint = f"{base_url}/classify"
+        self.classify_batch_endpoint = f"{base_url}/classify_batch"
         self.headers = headers
         self.cookies = cookies
         self.timeout = timeout
@@ -470,8 +472,34 @@ def classify(self, inputs: str) -> ClassifyResponse:
         if resp.status_code != 200:
             raise parse_error(resp.status_code, resp.json(), resp.headers if LORAX_DEBUG_MODE else None)
 
-        print(payload)
-        return ClassifyResponse(**payload)
+        return ClassifyResponse(entities=payload)
+    
+    def classify_batch(self, inputs: List[str]) -> List[ClassifyResponse]:
+        """
+        Given a list of inputs, run token classification on the text using the model
+
+        Args:
+            inputs (`List[str]`):
+                List of input texts
+        
+        Returns: 
+            List[Entities]: Entities found in the input text
+        """
+        request = BatchRequest(inputs=inputs)
+
+        resp = requests.post(
+            self.classify_batch_endpoint,
+            json=request.dict(by_alias=True),
+            headers=self.headers,
+            cookies=self.cookies,
+            timeout=self.timeout,
+        )
+
+        payload = resp.json()
+        if resp.status_code != 200:
+            raise parse_error(resp.status_code, resp.json(), resp.headers if LORAX_DEBUG_MODE else None)
+        
+        return [ClassifyResponse(entities=e) for e in payload]
 
 
 class AsyncClient:

clients/python/lorax/types.py

@@ -220,6 +220,28 @@ def valid_best_of_stream(cls, field_value, values):
         return field_value
 
 
+class BatchRequest(BaseModel):
+    # Prompt
+    inputs: List[str]
+    # Generation parameters
+    parameters: Optional[Parameters] = None
+    # Whether to stream output tokens
+    stream: bool = False
+
+    @field_validator("inputs")
+    def valid_input(cls, v):
+        if not v:
+            raise ValidationError("`inputs` cannot be empty")
+        return v
+
+    @field_validator("stream")
+    def valid_best_of_stream(cls, field_value, values):
+        parameters = values.data["parameters"]
+        if parameters is not None and parameters.best_of is not None and parameters.best_of > 1 and field_value:
+            raise ValidationError("`best_of` != 1 is not supported when `stream` == True")
+        return field_value
+
+
 # Decoder input tokens
 class InputToken(BaseModel):
     # Token ID from the model tokenizer