Update handler.py

handler.py

@@ -1,3 +1,4 @@
+
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -7,11 +8,7 @@ import io, base64, json, traceback, time
 from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
 import logging
 
-logging.basicConfig(
-    level=logging.DEBUG,
-    format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s",
-)
-
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s')
 
 class EndpointHandler:
     def __init__(self, model_dir, default_float16=True):
@@ -21,47 +18,46 @@ class EndpointHandler:
             self.processor = AutoProcessor.from_pretrained(
                 model_dir, trust_remote_code=True, torch_dtype="auto", device_map="auto"
             )
-            self.processor.tokenizer.padding_side = "left"
             self.model = AutoModelForCausalLM.from_pretrained(
-                model_dir, trust_remote_code=True, torch_dtype=dtype, device_map="auto"
+                model_dir,
+                trust_remote_code=True,
+                torch_dtype=dtype,
+                device_map="auto"
             )
         except Exception:
             logging.exception("Error en la inicialización del modelo")
             raise
 
-    def process_batch(
-        self,
-        prompts_list,
-        images_list,
-        images_config=None,
-        text_max_length=1535,
-        add_bos=True,
-    ):
+    def process_batch(self, prompts_list, images_list, images_config=None):
+        """
+        Ahora recibe una lista de prompts (strings) y la lista de imágenes,
+        en vez de un único 'prompt' replicado.
+        """
         try:
-            # Si se añade BOS, reducimos la longitud máxima en 1 para dejar espacio al token inicial.
-            token_max_length = text_max_length - 1 if add_bos else text_max_length
-
-            # Construimos los textos de entrada.
+            # Construimos el texto que va antes del prompt real
             batch_texts = [f"User: {p} Assistant:" for p in prompts_list]
 
-            # Tokenizamos con padding y truncamiento.
-            tokenized = self.processor.tokenizer(
-                batch_texts,
-                padding="max_length",
-                truncation=True,
-                max_length=token_max_length,
-                return_tensors="pt",
-            )
-            print(tokenized)
+            # Tokenizamos cada prompt por separado
+            tokens_list = [
+                self.processor.tokenizer.encode(" " + text, add_special_tokens=False)
+                for text in batch_texts
+            ]
+
+            # tokens_list = [
+            #     self.processor.tokenizer.encode(
+            #         " " + text,
+            #         add_special_tokens=False,
+            #         padding='longest',  # Asegura que todas las secuencias tengan la misma longitud
+            #         truncation=True,  # Opcional: trunca si hay una longitud máxima definida
+            #         max_length=1536  # Asegurar que no se pase del límite del modelo
+            #     )
+            #     for text in batch_texts
+            # ]
 
             outputs_list = []
             images_kwargs = {
-                "max_crops": images_config.get("max_crops", 12)
-                if images_config
-                else 12,
-                "overlap_margins": images_config.get("overlap_margins", [4, 4])
-                if images_config
-                else [4, 4],
+                "max_crops": images_config.get("max_crops", 12) if images_config else 12,
+                "overlap_margins": images_config.get("overlap_margins", [4, 4]) if images_config else [4, 4],
                 "base_image_input_size": [336, 336],
                 "image_token_length_w": 12,
                 "image_token_length_h": 12,
@@ -69,26 +65,21 @@ class EndpointHandler:
                 "image_padding_mask": True,
             }
 
-            # Preprocesamos cada imagen junto al prompt tokenizado.
+            # Para cada imagen y prompt, aplicamos el preprocesamiento multimodal
             for i in range(len(batch_texts)):
                 try:
-                    tokens = tokenized["input_ids"][i].tolist()
+                    tokens = tokens_list[i]
                     image = images_list[i].convert("RGB")
                     image = ImageOps.exif_transpose(image)
                     images_array = [np.array(image)]
-                    # Se espera que la secuencia final tenga 'text_max_length' tokens.
                     out = self.processor.image_processor.multimodal_preprocess(
                         images=images_array,
                         image_idx=[-1],
                         tokens=np.asarray(tokens).astype(np.int32),
-                        sequence_length=text_max_length,
-                        image_patch_token_id=self.processor.special_token_ids[
-                            "<im_patch>"
-                        ],
+                        sequence_length=1536,
+                        image_patch_token_id=self.processor.special_token_ids["<im_patch>"],
                         image_col_token_id=self.processor.special_token_ids["<im_col>"],
-                        image_start_token_id=self.processor.special_token_ids[
-                            "<im_start>"
-                        ],
+                        image_start_token_id=self.processor.special_token_ids["<im_start>"],
                         image_end_token_id=self.processor.special_token_ids["<im_end>"],
                         **images_kwargs,
                     )
@@ -97,63 +88,23 @@ class EndpointHandler:
                     logging.exception("Error procesando la imagen número %d", i)
                     raise
 
-            # Agrupamos las salidas en batch usando el token de padding.
-            pad_token_id = self.processor.tokenizer.pad_token_id
+            # Agrupamos las salidas en formato 'batch'
             batch_outputs = {}
             for key in outputs_list[0].keys():
                 try:
                     tensors = [torch.from_numpy(out[key]) for out in outputs_list]
                     batch_outputs[key] = torch.nn.utils.rnn.pad_sequence(
-                        tensors, batch_first=True, padding_value=pad_token_id
+                        tensors, batch_first=True, padding_value=self.processor.tokenizer.pad_token_id
                     )
                 except Exception:
-                    logging.exception(
-                        "Error al agrupar la key '%s' en outputs_list", key
-                    )
+                    logging.exception("Error al agrupar la key '%s' en outputs_list", key)
                     raise
 
-            # Calculamos la attention_mask a partir de input_ids.
-            attn_mask = (batch_outputs["input_ids"] != pad_token_id).long()
-
-            # Si se requiere, añadimos el token BOS al inicio.
-            if add_bos:
-                bos = (
-                    self.processor.tokenizer.bos_token_id
-                    or self.processor.tokenizer.eos_token_id
-                )
-                batch_outputs["input_ids"] = F.pad(
-                    batch_outputs["input_ids"], (1, 0), value=bos
-                )
-                attn_mask = F.pad(attn_mask, (1, 0), value=1)
-
-            # Si el modelo utiliza position_ids, calculamos position_ids y extendemos la attention_mask.
-            max_new_tokens_val = (
-                images_config.get("max_new_tokens", 0)
-                if images_config is not None
-                else 0
-            )
-            if self.model.config.use_position_ids and max_new_tokens_val > 0:
-                # Calculamos position_ids a partir de la atención (cumsum - 1, con mínimo 0).
-                position_ids = torch.clamp(
-                    torch.cumsum(attn_mask.to(torch.int32), dim=-1) - 1, min=0
-                )
-                # Calculamos append_last_valid_logits (la última posición válida en cada secuencia).
-                append_last_valid_logits = attn_mask.long().sum(dim=-1) - 1
-                # Extendemos la attention_mask a la derecha para incluir los nuevos tokens.
-                attn_mask = F.pad(attn_mask, (0, max_new_tokens_val), value=1)
-                # Guardamos estos valores en el batch.
-                batch_outputs["position_ids"] = position_ids
-                batch_outputs["append_last_valid_logits"] = append_last_valid_logits
-
-            # Asignamos la attention_mask calculada.
-            batch_outputs["attention_mask"] = attn_mask
+            # Ajuste para BOS token
+            bos = self.processor.tokenizer.bos_token_id or self.processor.tokenizer.eos_token_id
+            batch_outputs["input_ids"] = F.pad(batch_outputs["input_ids"], (1, 0), value=bos)
 
-            # Log para verificar la forma.
-            print(
-                f"[DEBUG] attention_mask.shape: {batch_outputs['attention_mask'].shape}"
-            )
-
-            # Ajuste de image_input_idx si existe.
+            # Ajustamos la posición de image_input_idx
             if "image_input_idx" in batch_outputs:
                 image_input_idx = batch_outputs["image_input_idx"]
                 batch_outputs["image_input_idx"] = torch.where(
@@ -183,11 +134,6 @@ class EndpointHandler:
             logging.exception("Error al acceder al campo 'inputs'")
             return {"error": "Error al acceder al campo 'inputs'."}
 
-        # Extraemos parámetros adicionales de configuración para pruebas (fuera de generation_config)
-        config_params = inputs_data.get("config", {})
-        text_max_length = config_params.get("text_max_length", 1535)
-        add_bos = config_params.get("add_bos", True)
-
         # Cargar imágenes y sus IDs
         images_list = []
         ids = []
@@ -206,15 +152,10 @@ class EndpointHandler:
                         images_list.append(image)
                         ids.append(image_id)
                     except Exception:
-                        logging.exception(
-                            "Error loading image with id %s",
-                            item.get("id", "desconocido"),
-                        )
+                        logging.exception("Error loading image with id %s", item.get("id", "desconocido"))
                         continue
             else:
-                return {
-                    "error": "Se requiere una lista de imágenes en 'inputs.images'."
-                }
+                return {"error": "Se requiere una lista de imágenes en 'inputs.images'."}
         except Exception:
             logging.exception("Error procesando la lista de imágenes")
             return {"error": "Error al procesar la lista de imágenes."}
@@ -223,12 +164,11 @@ class EndpointHandler:
         try:
             global_prompts_list = inputs_data.get("prompts", [])
             prompts_per_image = inputs_data.get("prompts_per_image", [])
+            # Diccionario: { image_id (str): [ {id, text}, {id, text}, ... ] }
             specific_prompts = {}
             for item in prompts_per_image:
                 if "id" in item and "prompts" in item:
-                    specific_prompts.setdefault(str(item["id"]), []).extend(
-                        item["prompts"]
-                    )
+                    specific_prompts.setdefault(str(item["id"]), []).extend(item["prompts"])
         except Exception:
             logging.exception("Error al construir el mapeo de prompts por imagen")
             return {"error": "Error al construir el mapeo de prompts por imagen."}
@@ -236,7 +176,7 @@ class EndpointHandler:
         # Preparamos la salida final
         final_results = {img_id: [] for img_id in ids}
 
-        # Configuración de generación (parámetros que se usan para el modelo)
+        # Configuración de generación
         try:
             batch_size = inputs_data.get("batch_size", len(images_list))
             generation_config = inputs_data.get("generation_config", {})
@@ -250,7 +190,7 @@ class EndpointHandler:
                 top_k=generation_config.get("top_k", 50),
                 length_penalty=generation_config.get("length_penalty", 1),
                 stop_strings="<|endoftext|>",
-                do_sample=True,
+                do_sample=True
             )
         except Exception:
             logging.exception("Error al configurar la generación")
@@ -260,6 +200,7 @@ class EndpointHandler:
         flattened = []
         try:
             for img, img_id in zip(images_list, ids):
+                # Si la imagen tiene prompts específicos, los usas. Si no, usas los globales
                 image_prompts = specific_prompts.get(str(img_id), global_prompts_list)
                 for p in image_prompts:
                     flattened.append((img, img_id, p["id"], p["text"]))
@@ -271,38 +212,16 @@ class EndpointHandler:
         print(f"[Info] Inicio de proceso por lotes sobre diccionario: {flattened}.")
         try:
             for start in range(0, len(flattened), batch_size):
-                chunk = flattened[start : start + batch_size]
-                # Registro de log para el lote actual (acortamos el prompt a 100 palabras)
-                batch_log = []
-                for item in chunk:
-                    photo_id = item[1]
-                    prompt_id = item[2]
-                    prompt_text = item[3]
-                    shortened = " ".join(prompt_text.split()[:100])
-                    batch_log.append(
-                        {
-                            "photo_id": photo_id,
-                            "prompt_id": prompt_id,
-                            "prompt_text": shortened,
-                        }
-                    )
-                logging.info(f"Lote {start // batch_size + 1}: {batch_log}")
-
+                chunk = flattened[start:start+batch_size]
+                # Extraemos imágenes y prompts
                 batch_imgs = [x[0] for x in chunk]
                 batch_img_ids = [x[1] for x in chunk]
                 batch_prompt_ids = [x[2] for x in chunk]
                 batch_prompt_texts = [x[3] for x in chunk]
 
-                inputs_batch = self.process_batch(
-                    batch_prompt_texts,
-                    batch_imgs,
-                    generation_config,
-                    text_max_length=text_max_length,
-                    add_bos=add_bos,
-                )
-                inputs_batch = {
-                    k: v.to(self.model.device) for k, v in inputs_batch.items()
-                }
+                # Preprocesamos
+                inputs_batch = self.process_batch(batch_prompt_texts, batch_imgs, generation_config)
+                inputs_batch = {k: v.to(self.model.device) for k, v in inputs_batch.items()}
 
                 if use_bfloat16 and "images" in inputs_batch:
                     inputs_batch["images"] = inputs_batch["images"].to(torch.bfloat16)
@@ -314,16 +233,19 @@ class EndpointHandler:
                         tokenizer=self.processor.tokenizer,
                     )
 
+                # Decodificamos
                 input_len = inputs_batch["input_ids"].shape[1]
                 generated_texts = self.processor.tokenizer.batch_decode(
                     outputs[:, input_len:], skip_special_tokens=True
                 )
 
+                # 3) Asignamos cada descripción generada a la imagen y prompt correctos
                 for idx, text in enumerate(generated_texts):
-                    final_results[batch_img_ids[idx]].append(
-                        {"id_prompt": batch_prompt_ids[idx], "description": text}
-                    )
-
+                    final_results[batch_img_ids[idx]].append({
+                        "id_prompt": batch_prompt_ids[idx],
+                        "description": text
+                    })
+                
                 torch.cuda.empty_cache()
 
         except Exception:
@@ -336,10 +258,12 @@ class EndpointHandler:
                 {"id": img_id, "descriptions": descs}
                 for img_id, descs in final_results.items()
             ]
-            print(
-                f"[DEBUG] Tiempo total de procesamiento: {time.time() - global_start_time:.2f} segundos."
-            )
+            print(f"[DEBUG] Tiempo total de procesamiento: {time.time() - global_start_time:.2f} segundos.")
             return combined_results
         except Exception:
             logging.exception("Error al combinar los resultados finales")
             return {"error": "Error al combinar los resultados finales."}
+
+
+
+