folder-classifier 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

folder_classifier/classifier.py

@@ -1,37 +1,25 @@
 from typing import Tuple
-
-import numpy as np
 import torch
 from transformers import pipeline
-
 from folder_classifier.dto import Listing
 
+
 classifier = None
 
-candidate_labels = ["legal_matter", "other"]
 
 def predict(listing: Listing) -> Tuple[str, float]:
     global classifier
     if classifier is None:
         classifier = pipeline(
-            "zero-shot-classification",
-            model="MoritzLaurer/ModernBERT-large-zeroshot-v2.0",
+            "text-classification",
+            model="/mnt/cluster_storage/models/corto-ai/ModernBERT-large-folder-classifier",
             torch_dtype=torch.bfloat16,
             device="cuda"
         )
     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
+    prediction = classifier(text)
+    predicted_label = prediction[0]["label"]
+    confidence = prediction[0]["score"]
+    return predicted_label, confidence
 
 

folder_classifier/dto.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from textwrap import dedent
-from typing import List, Union, Literal
+from typing import List, Union, Literal, Optional
 
 from pydantic import BaseModel, Field, confloat
 
@@ -24,7 +24,7 @@ class Folder(BaseModel):
     name: str
     type: Literal["folder"]
     # Discriminated union: 'type' field is used to select between File and Folder
-    items: List[Union[File, Folder]] = Field(default_factory=list)
+    items: Optional[List[Union[File, Folder]]] = Field(default_factory=list)
     model_config = {
         "json_schema_extra": {
             # Override the OpenAPI example to avoid the default 'string' entry
@@ -49,11 +49,6 @@ class Listing(BaseModel):
 Folder.model_rebuild()
 FolderClassificationRequest = Union[Folder, Listing]
 
-
-class ItemsRequest(BaseModel):
-    items: List[Union[File, Folder]]
-
-
 class FolderClassificationResponse(BaseModel):
     category: Literal["matter", "other"]
     confidence: confloat(ge=0.0, le=1.0)

{folder_classifier-0.2.1.dist-info → folder_classifier-0.2.2.dist-info}/METADATA

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

folder_classifier-0.2.2.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=YVKXvNAHny167H0Iv0GlRJEtRMAYwGea3cUVeGC4_sI,668
+folder_classifier/deploy.py,sha256=UQTbQjR_JX92Xo8L5EbK4nPg_VJmwWDmzMQSjP3-7iQ,412
+folder_classifier/dto.py,sha256=8ZhXO8q4OVIXxepOAtQjsYnoVL6yfqQc8WkY9bXEN5Q,1275
+folder_classifier/util.py,sha256=t-ma2suHovfNutogJb9jailRbRpIg4qv-zph3dHb2og,692
+folder_classifier-0.2.2.dist-info/METADATA,sha256=RHGa9U46FIqBKaoKDMlBw1L5yC4K-f-MZDwFd4qp5l0,392
+folder_classifier-0.2.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+folder_classifier-0.2.2.dist-info/top_level.txt,sha256=36ugc9pEbNQ-mnzz4Ot2WVjY3t_LzAN6XOCjDFP4p4k,18
+folder_classifier-0.2.2.dist-info/RECORD,,

folder_classifier-0.2.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-folder_classifier/__init__.py,sha256=k0YWZyUNe7myJiKeX0OaXtJ30_3EGE-vsZiAUbqa-3E,46
-folder_classifier/app.py,sha256=FqoBp_KQ3yIfoHTagdDTVsLbIj8luCn7en533Q870x8,1443
-folder_classifier/classifier.py,sha256=vt7LxIKpzGcB-74epsQKrEg6CfPuJnoRsQ-qIkO7r8o,1062
-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.1.dist-info/METADATA,sha256=8vaTYk8ZVnqcmZaNAs3AiC9iWXV4XopeuFlquj5CiBg,392
-folder_classifier-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-folder_classifier-0.2.1.dist-info/top_level.txt,sha256=36ugc9pEbNQ-mnzz4Ot2WVjY3t_LzAN6XOCjDFP4p4k,18
-folder_classifier-0.2.1.dist-info/RECORD,,