fix: serve not using defined model

Author: jianshen92

pdm.lock

@@ -38,7 +38,7 @@ dependencies = [
     "pydantic",
     "Pillow",
     "open-clip-torch",
-    "torch",
+    "torch==2.0.0",
     "typer",
 ]
 

pyproject.toml

@@ -43,7 +43,7 @@ class RankOutput(BaseItem):
     probabilities: List[List[float]]
     cosine_similarities: List[List[float]]
 
-bento_model = init_model()
+bento_model = init_model("__model_name__")
 logit_scale = np.exp(bento_model.info.metadata.get("logit_scale", 4.60517))
 
 clip_runner = get_clip_runner(bento_model)

src/clip_api_service/service.py src/clip_api_service/_service.pysrc/clip_api_service/service.py renamed to src/clip_api_service/_service.py

@@ -20,6 +20,19 @@ def build_bento(model_name: str | None = None, use_gpu: bool = False):
     if model_name:
         os.environ[MODEL_ENV_VAR_KEY] = model_name
 
+    # Generate a service file
+    # TODO: Refactor the code gen to a mode generalize function
+    # TODO: Use temporary file instead of writing to the source file 
+    src_service_file_path = os.path.join(os.path.dirname(__file__), "_service.py")
+    target_service_file_path = os.path.join(os.path.dirname(__file__), "service.py")
+
+    with open(src_service_file_path, 'r') as src_file, open(target_service_file_path, 'w') as target_file:
+        data = src_file.read().replace('__model_name__', model_name)
+        target_file.write(data)
+
     # Get parent directory path
     build_ctx = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
     bentoml.bentos.build_bentofile(bento_file, build_ctx=build_ctx)
+
+    # Remove the generated service file
+    os.remove(target_service_file_path)

src/clip_api_service/build.py

@@ -165,14 +165,14 @@ def __init__(self, bento_model: bentoml.Model):
     def encode_text(self, texts: list[str]) -> npt.NDArray:
         texts_encodings = self.tokenizer(texts)
         with torch.inference_mode():
-            text_embeddings = self.model.encode_text(texts_encodings)
+            text_embeddings = self.model.encode_text(texts_encodings.to(self.device))
             return text_embeddings.cpu().detach().numpy()
 
     @bentoml.Runnable.method(batchable=True)
     def encode_image(self, images: list[Image.Image]) -> npt.NDArray:
         image_encodings = torch.stack([self.processor(image) for image in images])
         with torch.inference_mode():
-            image_embeddings = self.model.encode_image(image_encodings)
+            image_embeddings = self.model.encode_image(image_encodings.to(self.device))
             return image_embeddings.cpu().detach().numpy()
 
 

src/clip_api_service/models/openclip.py

@@ -1,6 +1,6 @@
 import pprint
 
-from clip_api_service.service import svc
+from clip_api_service._service import svc
 
 input_sample = {
     "queries": [
@@ -28,7 +28,7 @@
 def test():
     import asyncio
 
-    from clip_api_service.service import RankInput
+    from clip_api_service._service import RankInput
 
     rank_input = RankInput.parse_obj(input_sample)