folder-classifier 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

folder_classifier/app.py

@@ -3,7 +3,10 @@ import logging
 from fastapi import FastAPI
 from ray import serve
 
-from folder_classifier.dto import ModelConfig, FolderClassificationResponse, FolderClassificationRequest
+from folder_classifier.dto import (ModelConfig, FolderClassificationResponse,
+                                   FolderClassificationRequest, Listing, Folder)
+from folder_classifier.util import flatten_folder
+from folder_classifier import classifier
 
 web_api = FastAPI(title=f"Folder Classifier API")
 
@@ -22,6 +25,7 @@ class FolderClassifierAPI:
 
     @web_api.post("/predict")
     async def predict(self, request: FolderClassificationRequest) -> FolderClassificationResponse:
-        result = ("matter", 0.9)  #await self.model_handle.remote(request)
+        listing = request if isinstance(request, Listing) else Listing(items=flatten_folder(request))
+        category, confidence = classifier.predict(listing)
         self.logger.info(f"Received request: {request}")
-        return FolderClassificationResponse(category=result[0], confidence=result[1])
+        return FolderClassificationResponse(category=category, confidence=confidence)

folder_classifier/classifier.py

@@ -0,0 +1,34 @@
+from typing import Tuple
+
+import numpy as np
+import torch
+from transformers import pipeline
+
+from folder_classifier.dto import Listing
+
+classifier = pipeline(
+    "zero-shot-classification",
+    model="MoritzLaurer/ModernBERT-large-zeroshot-v2.0",
+    torch_dtype=torch.bfloat16,
+    device="cuda"
+)
+
+candidate_labels = ["legal_matter", "other"]
+
+def predict(listing: Listing) -> Tuple[str, float]:
+    text = "\n".join(listing.items)
+    hypothesis_template = "This list of files is about {}"
+    prediction = classifier(
+        text,
+        candidate_labels,
+        hypothesis_template=hypothesis_template,
+        multi_label=False,
+    )
+    scores = np.array(prediction["scores"], dtype=float)
+    highest_ix = np.argmax(scores)
+    predicted_label = prediction["labels"][highest_ix]
+    confidence = float(scores[highest_ix])
+    prediction = "matter" if predicted_label == "legal_matter" else "other"
+    return prediction, confidence
+
+

folder_classifier/dto.py

@@ -25,25 +25,34 @@ class Folder(BaseModel):
     type: Literal["folder"]
     # Discriminated union: 'type' field is used to select between File and Folder
     items: List[Union[File, Folder]] = Field(default_factory=list)
-    # model_config = {
-    #     "json_schema_extra": {
-    #         # Override the OpenAPI example to avoid the default 'string' entry
-    #         "example": dedent("""{
-    #             "name": "string",
-    #             "type": "folder",
-    #             "items": [
-    #                 {
-    #                   "name": "string",
-    #                   "type": "file"
-    #                 }
-    #               ]
-    #             }""")
-    #     }
-    # }
+    model_config = {
+        "json_schema_extra": {
+            # Override the OpenAPI example to avoid the default 'string' entry
+            "example": dedent("""{
+                "name": "string",
+                "type": "folder",
+                "items": [
+                    {
+                      "name": "string",
+                      "type": "file"
+                    }
+                  ]
+                }""")
+        }
+    }
+
+
+class Listing(BaseModel):
+    items: List[str]
 
 
 Folder.model_rebuild()
-FolderClassificationRequest = Folder
+FolderClassificationRequest = Union[Folder, Listing]
+
+
+class ItemsRequest(BaseModel):
+    items: List[Union[File, Folder]]
+
 
 class FolderClassificationResponse(BaseModel):
     category: Literal["matter", "other"]

folder_classifier/util.py

@@ -0,0 +1,21 @@
+from typing import List
+
+from folder_classifier.dto import Folder
+
+
+def flatten_folder(folder: Folder, parent_path: str = "") -> List[str]:
+    """
+    Traverses a Folder and returns a list of file paths.
+    Each path is constructed by joining folder and file names with '/'.
+    """
+    paths: List[str] = []
+    # Build the path for the current folder
+    current_path = f"{parent_path}/{folder.name}" if parent_path else folder.name
+
+    for item in folder.items:
+        if item.type == "file":
+            paths.append(f"{current_path}/{item.name}")
+        else:
+            # Recursively flatten subfolders
+            paths.extend(flatten_folder(item, current_path))
+    return paths

{folder_classifier-0.1.1.dist-info → folder_classifier-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: folder-classifier
-Version: 0.1.1
+Version: 0.2.0
 Summary: Deploy folder classifier API to a Ray cluster
 Author: Crispin Almodovar
 Author-email: 

folder_classifier-0.2.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+folder_classifier/__init__.py,sha256=k0YWZyUNe7myJiKeX0OaXtJ30_3EGE-vsZiAUbqa-3E,46
+folder_classifier/app.py,sha256=FqoBp_KQ3yIfoHTagdDTVsLbIj8luCn7en533Q870x8,1443
+folder_classifier/classifier.py,sha256=zgkji_FfWILAvf6yN56doVF917Nx6lOrWhDw4hYoJ-k,947
+folder_classifier/deploy.py,sha256=UQTbQjR_JX92Xo8L5EbK4nPg_VJmwWDmzMQSjP3-7iQ,412
+folder_classifier/dto.py,sha256=WFWkYUVn5rrAowObBlMyacbgzHB30km0Y4vsAPjOULc,1326
+folder_classifier/util.py,sha256=t-ma2suHovfNutogJb9jailRbRpIg4qv-zph3dHb2og,692
+folder_classifier-0.2.0.dist-info/METADATA,sha256=Hfln_xbCUIhoV5CzZr0i2uKPXhMtILZLwMxknyOlc9Q,392
+folder_classifier-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+folder_classifier-0.2.0.dist-info/top_level.txt,sha256=36ugc9pEbNQ-mnzz4Ot2WVjY3t_LzAN6XOCjDFP4p4k,18
+folder_classifier-0.2.0.dist-info/RECORD,,

folder_classifier-0.1.1.dist-info/RECORD

@@ -1,8 +0,0 @@
-folder_classifier/__init__.py,sha256=k0YWZyUNe7myJiKeX0OaXtJ30_3EGE-vsZiAUbqa-3E,46
-folder_classifier/app.py,sha256=wMFvoPvk6coxX7t097h0EbdI4SHlT8hNS69gRZUAK-k,1212
-folder_classifier/deploy.py,sha256=UQTbQjR_JX92Xo8L5EbK4nPg_VJmwWDmzMQSjP3-7iQ,412
-folder_classifier/dto.py,sha256=CpFoQtiOP8-ZILkChBFRmNvP4oQZMI98Q5sgxldICLM,1220
-folder_classifier-0.1.1.dist-info/METADATA,sha256=lzNI0LRjdx4SZxIioRtBngWvu40kO401WIrnRx-Sdwg,392
-folder_classifier-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-folder_classifier-0.1.1.dist-info/top_level.txt,sha256=36ugc9pEbNQ-mnzz4Ot2WVjY3t_LzAN6XOCjDFP4p4k,18
-folder_classifier-0.1.1.dist-info/RECORD,,