Update handler.py

handler.py

@@ -10,6 +10,7 @@ import json
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from PIL import Image
 from transformers import AutoProcessor, LlavaForConditionalGeneration
 from transformers.utils import logging
@@ -77,8 +78,6 @@ class PixtralForRegression(nn.Module):
         am = attention_mask.to(last_h.device).long()
         ids = input_ids.to(last_h.device)
 
-
-
         if self.pooling == "mean_image_tokens":
             # Mean over all [IMG] placeholder tokens
             img_mask = (ids == self.image_token_id) & (am == 1)          # [B,L]
@@ -102,11 +101,6 @@ class PixtralForRegression(nn.Module):
             bsz = last_h.size(0)
             return last_h[torch.arange(bsz, device=last_h.device), idx]  # [B,H]
 
-
-
-        raise ValueError(f"Unknown pooling: {self.pooling}")
-
-
         raise ValueError(f"Unknown pooling: {self.pooling}")
 
     def forward(self, input_ids, attention_mask, pixel_values, **kwargs):
@@ -146,7 +140,7 @@ class PixtralForRegression(nn.Module):
             self._dbg = True
             img_mask = (input_ids == self.image_token_id) & (attention_mask == 1)
             print("IMG tokens per sample:", img_mask.sum(dim=1)[:4].tolist())
-        pooled = self._pool(last_h, attention_mask)
+        pooled = self._pool(last_h, attention_mask, input_ids)
         raw = self.reg_head(pooled.to(torch.float32)).squeeze(-1)
         preds = F.softplus(raw) + 1.0
         return {"logits": preds}
@@ -249,21 +243,6 @@ class EndpointHandler:
             raise ValueError("Missing 'subcat' (or 'sub_category') in 'inputs'.")
         return str(subcat).strip()
 
-    def _build_chat_text(self, prompt: str) -> str:
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "image"},
-                    {"type": "text", "text": prompt}
-                ],
-            }
-        ]
-        return self.processor.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            tokenize=False,
-        )
 
     def _build_regression_text(self, prompt: str) -> str:
         """
@@ -285,8 +264,8 @@ class EndpointHandler:
             messages,
             add_generation_prompt=False,
             tokenize=False,
-        ).rstrip()
-        return chat + "\n\nANSWER:"
+        )
+        return chat
 
     def __call__(self, data: Dict[str, Any]) -> Any:
         inputs = data.get("inputs", data)
@@ -301,7 +280,7 @@ class EndpointHandler:
         if not sub_category:
             raise ValueError("Missing 'sub_category' (or 'subcat') in 'inputs'.")
 
-        prompt = DEFAULT_PROMPT.format(SUB_CATEGORY=sub_category)
+        prompt = DEFAULT_PROMPT.format(SUB_CATEGORY=sub_category).rstrip() + "\n\nANSWER:"
 
         image = self._decode_image(image_b64)
         image = self._resize_max_side(image, max_side=int(inputs.get("max_side", self.cfg.max_side)))