removed mp

handler.py

@@ -2,13 +2,12 @@ import logging
 from typing import List, Dict
 
 import torch.multiprocessing as mp
-from pydantic import BaseModel
 from sentence_transformers import CrossEncoder
 
 import config
 
 # Used for CUDA multiprocessing
-mp.set_start_method('spawn', force=True)
+# mp.set_start_method('spawn', force=True)
 
 formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 console_handler = logging.StreamHandler()
@@ -19,21 +18,6 @@ logger.addHandler(console_handler)
 logger.setLevel(logging.DEBUG)
 
 
-class DocsToRerank(BaseModel):
-    query: str
-    documents: List[str]
-
-
-class Payload(BaseModel):
-    inputs: DocsToRerank
-    parameters: Dict[str, str] = None
-
-
-class Rank(BaseModel):
-    corpus_id: int
-    score: float
-
-
 class EndpointHandler:
     def __init__(self, path):
         self.top_k = config.CROSSENCODER_TOP_K
@@ -55,7 +39,7 @@ class EndpointHandler:
             inputs['query'],
             inputs['documents'],
             top_k=self.top_k,
-            num_workers=self.max_workers
+            # num_workers=self.max_workers
         )
         logger.info(f'New ranks: {ranks}')
 

test_endpoint_handler.py

@@ -1,6 +1,6 @@
 import pytest
 
-from handler import EndpointHandler, Payload
+from handler import EndpointHandler
 
 
 @pytest.mark.parametrize('payload', [