Update inference_wrapper.py

inference_wrapper.py

@@ -1,4 +1,3 @@
-
 import numpy as np
 import joblib
 import torch
@@ -9,6 +8,8 @@ class FinancialFilingClassifier:
     def __init__(self, model_dir):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Loading Jina Encoder on {self.device}...")
+        
+        # Jina-v3 handles Flash Attention internally if installed
         self.encoder = AutoModel.from_pretrained(
             "jinaai/jina-embeddings-v3", 
             trust_remote_code=True, 
@@ -21,13 +22,40 @@ class FinancialFilingClassifier:
         self.specialists = {}
         self.model_dir = model_dir
 
-    def _get_vector(self, text):
-        log_len = np.log1p(len(str(text)))
+    def _get_vector(self, text, doc_length=None):
+        """
+        Generates the feature vector: [Embedding (1024) + Log_Length (1)]
+        
+        Args:
+            text (str): The document text content.
+            doc_length (int, optional): The real length of the document. 
+                                      If None, defaults to len(text).
+        """
+        # 1. Feature Engineering: Log Length
+        # We use the explicitly passed length if available. This allows
+        # the model to know a document is massive (e.g. Annual Report) 
+        # even if we only run inference on the first 8k tokens.
+        val = int(doc_length) if doc_length is not None else len(str(text))
+        log_len = np.log1p(val)
+
         with torch.no_grad():
-            vec = self.encoder.encode([text], task="classification", max_length=8192)
+            # 2. Embedding Generation
+            # SAFETY MARGIN: We cap at 8100 (instead of 8192) to prevent 
+            # "Rotary Embedding" off-by-one crashes in Flash Attention.
+            vec = self.encoder.encode([text], task="classification", max_length=8100)
+            
+            # 3. XGBoost Compatibility
+            # XGBoost requires CPU-bound Numpy arrays.
+            if isinstance(vec, torch.Tensor):
+                vec = vec.cpu().numpy()
+            elif isinstance(vec, list):
+                vec = np.array(vec)
+
+        # 4. Concatenate: [Embedding vector, Log_Length]
         return np.hstack([vec, [[log_len]]])
 
     def _load_specialist(self, category):
+        """Lazy loads specialist models to save RAM until needed."""
         safe_name = category.replace(" ", "_").replace("&", "and").replace("/", "_")
         if safe_name not in self.specialists:
             try:
@@ -38,15 +66,28 @@ class FinancialFilingClassifier:
                 return None
         return self.specialists[safe_name]
 
-    def predict(self, text):
-        vector = self._get_vector(text)
+    def predict(self, text, doc_length=None):
+        """
+        Predicts the category and type of a financial document.
+        
+        Args:
+            text (str): The document text.
+            doc_length (int, optional): The true character length of the document.
+                                      Recommended for highest accuracy.
+        """
+        vector = self._get_vector(text, doc_length=doc_length)
+        
+        # 1. Router Prediction (General Category)
         router_probs = self.router.predict_proba(vector)[0]
+        # We look at the top 2 candidates to handle ambiguous edge cases
         top_indices = np.argsort(router_probs)[::-1][:2]
         
         candidates = []
         for idx in top_indices:
             category = self.router_le.classes_[idx]
             router_conf = router_probs[idx]
+            
+            # 2. Specialist Prediction (Specific Type)
             specialist = self._load_specialist(category)
             
             if specialist:
@@ -55,9 +96,21 @@ class FinancialFilingClassifier:
                 best_idx = np.argmax(spec_probs)
                 label = le.classes_[best_idx]
                 spec_conf = spec_probs[best_idx]
+                
+                # Combine Confidence Scores (Geometric Mean)
                 combined_score = np.sqrt(router_conf * spec_conf)
-                candidates.append({"category": category, "label": label, "score": float(combined_score)})
+                candidates.append({
+                    "category": category, 
+                    "label": label, 
+                    "score": float(combined_score)
+                })
             else:
-                candidates.append({"category": category, "label": category, "score": float(router_conf)})
+                # Fallback if no specialist exists for this category
+                candidates.append({
+                    "category": category, 
+                    "label": category, 
+                    "score": float(router_conf)
+                })
         
-        return max(candidates, key=lambda x: x['score'])
+        # Return the highest scoring candidate
+        return max(candidates, key=lambda x: x['score'])