Update handler.py

handler.py

@@ -1,184 +1,426 @@
-from __future__ import annotations
-
-import base64
-from dataclasses import dataclass
-from io import BytesIO
-from typing import Any, Dict, Optional, List
-
-import torch
-from PIL import Image
-from transformers import AutoProcessor, LlavaForConditionalGeneration
-from transformers.utils import logging
-
-
-logger = logging.get_logger(__name__)
-logging.set_verbosity_info()
-
-
-BASE_MODEL_ID = "mistral-community/pixtral-12b"
-
-
-# Prompt par défaut (tu peux l’ajuster ici)
-DEFAULT_PROMPT = (
-    "Here is a photo showing some food waste. "
-    "Identify each type of food item and the corresponding weight in grams. "
-    "Reply like: Milk, 120g; Coffee, 45g. "
-    "Do not add any explanation, no extra text."
-)
-
-
-@dataclass
-class GenerationConfig:
-
-    max_new_tokens: int = 64
-    temperature: float = 0.0
-    no_repeat_ngram_size: int = 6
-    repetition_penalty: float = 1.1
-
-
-class EndpointHandler:
-    
-
-    def __init__(self, path: str = ".") -> None:
-        """
-        Initializes the model and processor from the `path` directory,
-        which contains the merged weights (pixtral-12b-foodwaste-merged).
-        """
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info("Initializing EndpointHandler on device: %s", self.device)
-
-        self.processor = AutoProcessor.from_pretrained(
-            BASE_MODEL_ID,
-            trust_remote_code=True,
-        )
-
-        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-        self.model = LlavaForConditionalGeneration.from_pretrained(
-            BASE_MODEL_ID,
-            torch_dtype=dtype,
-            low_cpu_mem_usage=True,
-            device_map={"": self.device},
-            trust_remote_code=True,
-        )
-        self.model.eval()
-        logger.info("Model and processor successfully loaded from '%s'.", path)
-
-        # pad token management
-        tokenizer = getattr(self.processor, "tokenizer", None)
-        if tokenizer is not None and tokenizer.pad_token_id is None:
-            tokenizer.pad_token = tokenizer.eos_token
-            tokenizer.pad_token_id = tokenizer.eos_token_id
-
-        # Preparation of EOS/PAD IDs for generate
-        eos_candidates: List[int] = []
-        if self.model.config.eos_token_id is not None:
-            eos_candidates.append(self.model.config.eos_token_id)
-        if tokenizer is not None and tokenizer.eos_token_id is not None:
-            eos_candidates.append(tokenizer.eos_token_id)
-
-        self.eos_token_ids: List[int] = list({i for i in eos_candidates})
-        if not self.eos_token_ids:
-            raise ValueError("No EOS token id found on model or tokenizer.")
-
-        pad_id: Optional[int] = getattr(self.model.config, "pad_token_id", None)
-        if pad_id is None and tokenizer is not None:
-            pad_id = tokenizer.pad_token_id
-        if pad_id is None:
-            pad_id = self.eos_token_ids[0]
-
-        self.pad_token_id: int = pad_id
-
-        self.gen_config = GenerationConfig()
-        logger.info(
-            "Generation config: max_new_tokens=%d, temperature=%.3f",
-            self.gen_config.max_new_tokens,
-            self.gen_config.temperature,
-        )
-
-    
-    @staticmethod
-    def _decode_image(image_b64: str) -> Image.Image:
-        
-        try:
-            img_bytes = base64.b64decode(image_b64)
-            img = Image.open(BytesIO(img_bytes)).convert("RGB")
-            return img
-        except Exception as exc:  # pragma: no cover - log production
-            raise ValueError(f"Could not decode base64 image: {exc}") from exc
-
-    def _build_chat_text(self, prompt: str) -> str:
-        
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt},
-                    {"type": "image"},
-                ],
-            }
-        ]
-
-        chat_text = self.processor.apply_chat_template(
-            messages,
-            add_generation_prompt=True,
-            tokenize=False,
-        )
-        return chat_text
-
-    
-    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        
-        inputs = data.get("inputs", data)
-        
-        prompt: str = inputs.get("prompt") or DEFAULT_PROMPT
-
-        image_b64: Optional[str] = inputs.get("image")
-        if image_b64 is None:
-            raise ValueError("Missing 'image' field (base64-encoded) in 'inputs'.")
-
-        image = self._decode_image(image_b64)
-
-        max_new_tokens = int(inputs.get("max_new_tokens", self.gen_config.max_new_tokens))
-        temperature = float(inputs.get("temperature", self.gen_config.temperature))
-
-        logger.info(
-            "Received request: max_new_tokens=%d, temperature=%.3f",
-            max_new_tokens,
-            temperature,
-        )
-
-        chat_text = self._build_chat_text(prompt)
-
-        enc = self.processor(
-            text=[chat_text],
-            images=[image],
-            return_tensors="pt",
-            truncation=False,
-        )
-        enc = {k: v.to(self.device) for k, v in enc.items()}
-        if "pixel_values" in enc:
-            enc["pixel_values"] = enc["pixel_values"].to(self.device, dtype=self.model.dtype)
-
-        gen_kwargs: Dict[str, Any] = {
-            "max_new_tokens": max_new_tokens,
-            "do_sample": temperature > 0.0,
-            "eos_token_id": self.eos_token_ids,
-            "pad_token_id": self.pad_token_id,
-            "no_repeat_ngram_size": self.gen_config.no_repeat_ngram_size,
-            "repetition_penalty": self.gen_config.repetition_penalty,
-        }
-        if temperature > 0.0:
-            gen_kwargs["temperature"] = temperature
-
-        with torch.inference_mode():
-            output_ids = self.model.generate(**enc, **gen_kwargs)
-
-        generated_only = output_ids[:, enc["input_ids"].shape[1]:]
-        generated_text = self.processor.batch_decode(
-            generated_only,
-            skip_special_tokens=True,
-        )[0].strip()
-
-        logger.info("Generated text: %s", generated_text)
-
-        return {"generated_text": generated_text}
+from __future__ import annotations
+
+import base64
+from curses import raw
+import re
+import unicodedata
+from dataclasses import dataclass
+from io import BytesIO
+from typing import Any, Dict, Optional, List, Set
+
+import torch
+from PIL import Image
+from transformers import AutoProcessor, LlavaForConditionalGeneration
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+logging.set_verbosity_info()
+
+DEFAULT_PROMPT = (
+    "Here is a picture showing some food waste.\n\n"
+    "Task: Provide a list of the food waste items visible in the picture.\n"
+    "For each item, output EXACTLY one CATEGORY chosen from the CATEGORY DICTIONARY below.\n\n"
+    "Rules:\n"
+    "- Output must be a single line (no extra text, no markdown).\n"
+    "- Use the exact spelling of CATEGORY as listed.\n"
+    "- If you are unsure, choose the closest broader category (e.g., 'fruit', 'vegetable', 'food_waste').\n\n"
+    "CATEGORY DICTIONARY:\n"
+    "1. cucumber\n"
+    "2. tomato\n"
+    "3. cherry tomato\n"
+    "4. carrot\n"
+    "5. lettuce\n"
+    "6. bell_pepper\n"
+    "7. cabbage\n"
+    "8. cauliflower\n"
+    "9. broccoli\n"
+    "10. onion\n"
+    "11. herbs\n"
+    "12. vegetable_mix\n"
+    "13. vegetable\n"
+    "14. potato\n"
+    "15. potato_product\n"
+    "16. potato_based\n"
+    "17. banana\n"
+    "18. apple\n"
+    "19. mandarin\n"
+    "20. orange\n"
+    "21. lemon\n"
+    "22. grape\n"
+    "23. plum\n"
+    "24. nectarine\n"
+    "25. dried fruits\n"
+    "26. jam\n"
+    "27. cherry\n"
+    "28. fruit\n"
+    "29. blueberry\n"
+    "30. strawberry\n"
+    "31. raspberry\n"
+    "32. currant\n"
+    "33. lingonberry\n"
+    "34. berries\n"
+    "35. rice\n"
+    "36. pasta\n"
+    "37. rice_or_pasta\n"
+    "38. dark_bread\n"
+    "39. light_bread\n"
+    "40. bread\n"
+    "41. oatmeal\n"
+    "42. semolina_porridge\n"
+    "43. rice_porridge\n"
+    "44. whipped_porridge\n"
+    "45. porridge\n"
+    "46. savory_pastry\n"
+    "47. sweet_pastry\n"
+    "48. biscuit\n"
+    "49. baked_goods\n"
+    "50. flour\n"
+    "51. flakes\n"
+    "52. muesli\n"
+    "53. cereal\n"
+    "54. grain_products\n"
+    "55. milk\n"
+    "56. plant_based_milk\n"
+    "57. children_milk\n"
+    "58. dairy_product_milk\n"
+    "59. block_cheese\n"
+    "60. sliced_cheese\n"
+    "61. fresh_cheese\n"
+    "62. soft_cheese\n"
+    "63. plant_cheese\n"
+    "64. cheese\n"
+    "65. yoghurt\n"
+    "66. buttermilk\n"
+    "67. quark\n"
+    "68. fermented_milk\n"
+    "69. cream\n"
+    "70. ice_cream\n"
+    "71. plant_based_yogurt\n"
+    "72. plant_based_cream\n"
+    "73. plant_based_ice_cream\n"
+    "74. dairy_product\n"
+    "75. egg\n"
+    "76. eggs\n"
+    "77. beef_steak\n"
+    "78. minced_beef\n"
+    "79. cold_cut_beef\n"
+    "80. beef_sausage\n"
+    "81. beef_frankfurter\n"
+    "82. beef_product\n"
+    "83. beef\n"
+    "84. pork_steak\n"
+    "85. minced_pork\n"
+    "86. cold_cut_pork\n"
+    "87. pork_sausage\n"
+    "88. pork_frankfurter\n"
+    "89. pork_product\n"
+    "90. pork\n"
+    "91. chicken\n"
+    "92. minced_chicken\n"
+    "93. cold_cut_chicken\n"
+    "94. chicken_sausage\n"
+    "95. chicken_frankfurter\n"
+    "96. chicken_product\n"
+    "97. steak\n"
+    "98. minced_meat\n"
+    "99. cold_cut_meat\n"
+    "100. sausage\n"
+    "101. frankfurter\n"
+    "102. meat_product\n"
+    "103. plant_protein\n"
+    "104. meat\n"
+    "105. fish_fillet\n"
+    "106. fish_strips\n"
+    "107. fish_product\n"
+    "108. shellfish\n"
+    "109. fish\n"
+    "110. meat_soup\n"
+    "111. meat_stew\n"
+    "112. meat_sauce\n"
+    "113. meat_pizza\n"
+    "114. meat_hamburger\n"
+    "115. meat_salad\n"
+    "116. meat_dish\n"
+    "117. fish_soup\n"
+    "118. fish_stew\n"
+    "119. fish_sauce\n"
+    "120. fish_pizza\n"
+    "121. fish_hamburger\n"
+    "122. fish_salad\n"
+    "123. fish_dish\n"
+    "124. vegetarian_soup\n"
+    "125. vegetarian_stew\n"
+    "126. vegetarian_sauce\n"
+    "127. vegetarian_pizza\n"
+    "128. vegetarian_hamburger\n"
+    "129. vegetarian_salad\n"
+    "130. vegetarian_dish\n"
+    "131. pastry\n"
+    "132. sweet_soup\n"
+    "133. quark\n"
+    "134. ice_cream\n"
+    "135. crepes\n"
+    "136. jam\n"
+    "137. dessert\n"
+    "138. chocolate\n"
+    "139. candy\n"
+    "140. sweet\n"
+    "141. popcorn\n"
+    "142. potato_chip\n"
+    "143. snack\n"
+    "144. sauce_paste\n"
+    "145. spice\n"
+    "146. sauce_seasoning\n"
+    "147. butter\n"
+    "148. oil\n"
+    "149. fat_oil\n"
+    "150. juice\n"
+    "151. soda\n"
+    "152. soft_drink\n"
+    "153. coffee\n"
+    "154. tea\n"
+    "155. cocoa\n"
+    "156. hot_beverage\n"
+    "157. beer\n"
+    "158. cider\n"
+    "159. alcoholic_long_drink\n"
+    "160. wine\n"
+    "161. strong_alcoholic_beverage\n"
+    "162. alcoholic_beverage\n"
+    "163. sugar\n"
+    "164. honey\n"
+    "165. sirup\n"
+    "166. sugar_honey_syrup\n"
+    "167. food_waste\n\n"
+    "OUTPUT FORMAT (single line):\n"
+    "CATEGORY1,CATEGORY2,CATEGORY3\n"
+)
+
+
+@dataclass
+class GenerationConfig:
+
+    max_new_tokens: int = 256
+    temperature: float = 0.0
+    no_repeat_ngram_size: int = 6
+    repetition_penalty: float = 1.1
+    max_length: int = 4096
+    max_side: int = 512
+
+
+_ALLOWED_RE = re.compile(r"[^a-z0-9_\(\);,\|\/\-\s\?\.\:]")
+
+ITEM_RE = re.compile(r"\(\s*(.*?)\s*\)\s*\|\s*(\d+)", flags=re.DOTALL)
+
+def _clean_text(s: str) -> str:
+    s = unicodedata.normalize("NFKC", s).replace("\u00A0", " ")
+    s = re.sub(r"[\u200B-\u200F]", "", s)
+    s = re.sub(r"\s+", " ", s).strip()
+    return s
+
+
+_CATEGORY_LINE_RE = re.compile(r"^\s*\d+\.\s*(.+?)\s*$", flags=re.MULTILINE)
+
+def _extract_categories_from_prompt(prompt: str) -> Set[str]:
+    cats = {m.group(1).strip().lower() for m in _CATEGORY_LINE_RE.finditer(prompt)}
+    return {c for c in cats if c}
+
+def _clean_model_output(s: str) -> str:
+    s = _clean_text(s).lower()
+    s = _ALLOWED_RE.sub("", s)
+    return s
+
+def _parse_and_validate_categories(raw: str, allowed: Set[str]) -> List[str]:
+    s = _clean_model_output(raw)
+    parts = [p.strip() for p in s.split(",") if p.strip()]
+    out: List[str] = []
+    seen: Set[str] = set()
+    for p in parts:
+        if p in allowed and p not in seen:
+            out.append(p)
+            seen.add(p)
+    return out
+
+
+class EndpointHandler:
+    
+
+    def __init__(self, path: str = ".") -> None:
+        """
+        Initializes the model and processor from the `path` directory,
+        which contains the merged weights (pixtral-12b-foodwaste-merged).
+        """
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info("Initializing EndpointHandler on device: %s", self.device)
+        self.processor = AutoProcessor.from_pretrained(
+            path,
+            trust_remote_code=True,
+        )
+
+        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+        self.model = LlavaForConditionalGeneration.from_pretrained(
+            path,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+            device_map={"": self.device},
+            trust_remote_code=True,
+        )
+        self.model.eval()
+        logger.info("Model and processor successfully loaded from '%s'.", path)
+
+        # pad token management
+        tokenizer = getattr(self.processor, "tokenizer", None)
+        if tokenizer is not None and tokenizer.pad_token_id is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+
+        # Preparation of EOS/PAD IDs for generate
+        eos_candidates: List[int] = []
+        if self.model.config.eos_token_id is not None:
+            eos_candidates.append(self.model.config.eos_token_id)
+        if tokenizer is not None and tokenizer.eos_token_id is not None:
+            eos_candidates.append(tokenizer.eos_token_id)
+
+        self.eos_token_ids: List[int] = list({i for i in eos_candidates})
+        if not self.eos_token_ids:
+            raise ValueError("No EOS token id found on model or tokenizer.")
+
+        pad_id: Optional[int] = getattr(self.model.config, "pad_token_id", None)
+        if pad_id is None and tokenizer is not None:
+            pad_id = tokenizer.pad_token_id
+        if pad_id is None:
+            pad_id = self.eos_token_ids[0]
+
+        self.pad_token_id: int = pad_id
+
+        self.gen_config = GenerationConfig()
+        logger.info(
+            "Generation config: max_new_tokens=%d, temperature=%.3f",
+            self.gen_config.max_new_tokens,
+            self.gen_config.temperature,
+        )
+        
+        self.default_allowed_categories: Set[str] = _extract_categories_from_prompt(DEFAULT_PROMPT)
+        logger.info("Extracted %d categories from DEFAULT_PROMPT.", len(self.default_allowed_categories))
+ 
+
+    
+    @staticmethod
+    def _decode_image(image_b64: str) -> Image.Image:
+        
+        try:
+            img_bytes = base64.b64decode(image_b64)
+            img = Image.open(BytesIO(img_bytes)).convert("RGB")
+            return img
+        except Exception as exc:  # pragma: no cover - log production
+            raise ValueError(f"Could not decode base64 image: {exc}") from exc
+
+    @staticmethod
+    def _resize_max_side(img: Image.Image, max_side: int) -> Image.Image:
+        w, h = img.size
+        m = max(w, h)
+        if m <= max_side:
+            return img
+        scale = max_side / m
+        # LANCZOS = better downscale
+        return img.resize((int(w * scale), int(h * scale)), resample=Image.Resampling.LANCZOS)
+
+    def _build_chat_text(self, prompt: str) -> str:
+        
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {"type": "image"},
+                ],
+            }
+        ]
+
+        chat_text = self.processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=False,
+        )
+        return chat_text
+
+    
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        inputs = data.get("inputs", data)
+        debug = bool(inputs.get("debug", False))
+    
+        prompt: str = inputs.get("prompt") or DEFAULT_PROMPT
+        allowed_categories = _extract_categories_from_prompt(prompt) or self.default_allowed_categories
+        image_b64: Optional[str] = inputs.get("image")
+        if not image_b64:
+            raise ValueError("Missing 'image' field (base64-encoded) in 'inputs'.")
+    
+        image = self._decode_image(image_b64)
+    
+        max_length = int(inputs.get("max_length", self.gen_config.max_length))
+        max_side = int(inputs.get("max_side", self.gen_config.max_side))
+        image = self._resize_max_side(image, max_side=max_side)
+    
+        max_new_tokens = int(inputs.get("max_new_tokens", self.gen_config.max_new_tokens))
+        temperature = float(inputs.get("temperature", self.gen_config.temperature))
+    
+        chat_text = self._build_chat_text(prompt)
+    
+        enc = self.processor(
+            text=[chat_text],
+            images=[image],
+            return_tensors="pt",
+            truncation=True,
+            max_length=max_length,
+            padding=False,  # important for correct prompt_len
+        )
+    
+        prompt_len = int(enc["input_ids"].shape[1])
+        tokens_left = max(1, max_length - prompt_len)
+        max_new_tokens = min(max_new_tokens, tokens_left)
+    
+        if debug:
+            logger.info("===== TOKEN BUDGET DEBUG (inference) =====")
+            logger.info("img size: %s", getattr(image, "size", None))
+            logger.info("max_length: %d", max_length)
+            logger.info("prompt_len: %d", prompt_len)
+            logger.info("tokens_left_for_answer(approx): %d", tokens_left)
+            tok = getattr(self.processor, "tokenizer", None)
+            if tok is not None:
+                ids = enc["input_ids"][0].tolist()
+                logger.info("[prompt head]\n%s", tok.decode(ids[:120], skip_special_tokens=False))
+                logger.info("[prompt tail]\n%s", tok.decode(ids[-120:], skip_special_tokens=False))
+            logger.info("=========================================")
+    
+        enc = {k: v.to(self.device) for k, v in enc.items()}
+        if "pixel_values" in enc:
+            enc["pixel_values"] = enc["pixel_values"].to(self.device, dtype=self.model.dtype)
+    
+        gen_kwargs: Dict[str, Any] = {
+            "max_new_tokens": max_new_tokens,
+            "do_sample": temperature > 0.0,
+            "eos_token_id": self.eos_token_ids,
+            "pad_token_id": self.pad_token_id,
+            "no_repeat_ngram_size": self.gen_config.no_repeat_ngram_size,
+            "repetition_penalty": self.gen_config.repetition_penalty,
+        }
+        if temperature > 0.0:
+            gen_kwargs["temperature"] = temperature
+    
+        with torch.inference_mode():
+            output_ids = self.model.generate(**enc, **gen_kwargs)
+    
+        generated_only = output_ids[:, enc["input_ids"].shape[1]:]
+        generated_text = self.processor.batch_decode(
+            generated_only,
+            skip_special_tokens=True,
+        )[0].strip()
+    
+        cats = _parse_and_validate_categories(generated_text, allowed_categories)
+        if not cats:
+            cats = ["food_waste"] if "food_waste" in allowed_categories else []
+        generated_text = ",".join(cats)
+
+        logger.info("Generated text: %s", generated_text)
+        return {"generated_text": generated_text, "categories": cats}