Update modules/utils.py

modules/utils.py

@@ -1,21 +1,81 @@
-def handle_long_text(text, model, tokenizer, max_length=2048, stride=128):
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import torch
+from tqdm.auto import tqdm
+
+
+def handle_long_text(
+    input_text: str,
+    model: AutoModelForSeq2SeqLM,
+    tokenizer: AutoTokenizer,
+    max_length: int = 128,
+    stride: int = 128,
+    batch_length: int = 2048,
+    min_batch_length: int = 512,
+    **generate_kwargs,
+) -> str:
+    """
+    Maneja textos largos dividiéndolos en segmentos y generando resúmenes para cada uno.
+
+    Args:
+        input_text (str): Texto completo a resumir.
+        model: Modelo de resumen abstractivo.
+        tokenizer: Tokenizador asociado al modelo.
+        max_length (int): Longitud máxima del resumen generado por segmento.
+        stride (int): Cantidad de tokens que se superponen entre segmentos.
+        batch_length (int): Longitud máxima de tokens por segmento.
+        min_batch_length (int): Longitud mínima permitida por segmento.
+        generate_kwargs: Parámetros adicionales para el modelo de generación.
+
+    Returns:
+        str: Resumen final concatenado de todos los segmentos.
+    """
+    # Validar parámetros de longitud
+    if batch_length < min_batch_length:
+        batch_length = min_batch_length
+
+    # Tokenizar texto completo en segmentos
     encoded_input = tokenizer(
-        text,
-        max_length=max_length,
-        stride=stride,
+        input_text,
+        return_tensors="pt",
+        max_length=batch_length,
         truncation=True,
+        stride=stride,
         return_overflowing_tokens=True,
-        return_tensors="pt",
+        add_special_tokens=True,
     )
+    
+    # Obtener IDs y máscaras de atención
+    input_ids = encoded_input["input_ids"]
+    attention_masks = encoded_input["attention_mask"]
+
+    # Progresión para múltiples segmentos
     summaries = []
-    for input_ids, attention_mask in zip(
-        encoded_input.input_ids, encoded_input.attention_mask
-    ):
-        output = model.generate(
-            input_ids.to(model.device),
-            attention_mask=attention_mask.to(model.device),
-            max_length=128,
+    pbar = tqdm(total=len(input_ids), desc="Procesando segmentos")
+
+    for ids, mask in zip(input_ids, attention_masks):
+        # Enviar al dispositivo correcto (CPU/GPU)
+        ids = ids.unsqueeze(0).to(model.device)
+        mask = mask.unsqueeze(0).to(model.device)
+
+        # Generar resumen para el segmento actual
+        outputs = model.generate(
+            input_ids=ids,
+            attention_mask=mask,
+            max_length=max_length,
+            no_repeat_ngram_size=3,
             num_beams=4,
+            early_stopping=True,
+            **generate_kwargs,
+        )
+        # Decodificar resumen generado
+        summary = tokenizer.decode(
+            outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True
         )
-        summaries.append(tokenizer.decode(output[0], skip_special_tokens=True))
-    return " ".join(summaries)
+        summaries.append(summary)
+        pbar.update()
+
+    pbar.close()
+
+    # Concatenar resúmenes y devolver el texto final
+    final_summary = " ".join(summaries)
+    return final_summary