Update app.py

app.py

@@ -170,51 +170,40 @@ class DebiasingSFTTrainer(SFTTrainer):
         return (loss, outputs) if return_outputs else loss
 
 @spaces.GPU()
-class DeduplicatedIterableDataset(IterableDataset):
-    def __init__(self, dataset, text_col, method, threshold=0.85, num_perm=128):
-        super().__init__(ex_iterable=iter([]))
-        self.dataset = dataset
-        self.text_col = text_col
-        self.method = method
-        self.threshold = threshold
-        self.num_perm = num_perm
-        if hasattr(dataset, '_info'):
-            self._info = dataset._info
-        elif hasattr(dataset, 'info'):
-            self._info = dataset.info
-
-    def __iter__(self):
-        if self.method == 'Exacta':
-            return self._exact_iter()
-        elif self.method == 'Semántica (MinHash)':
-            return self._minhash_iter()
-        else:
-            return iter(self.dataset)
-
-    def _exact_iter(self):
-        seen_texts = set()
-        for example in self.dataset:
-            text = example.get(self.text_col, "")
-            if text and isinstance(text, str):
-                if text not in seen_texts:
-                    seen_texts.add(text)
+def _create_deduplicated_iterable_dataset(dataset, text_col, method, threshold=0.85, num_perm=128):
+    def gen():
+        if method == 'Exacta':
+            seen_texts = set()
+            for example in dataset:
+                text = example.get(text_col, "")
+                if text and isinstance(text, str):
+                    if text not in seen_texts:
+                        seen_texts.add(text)
+                        yield example
+                else:
                     yield example
-            else:
-                yield example
-
-    def _minhash_iter(self):
-        lsh = MinHashLSH(threshold=self.threshold, num_perm=self.num_perm)
-        for i, example in enumerate(self.dataset):
-            text = example.get(self.text_col, "")
-            if text and isinstance(text, str) and text.strip():
-                m = MinHash(num_perm=self.num_perm)
-                for d in text.split():
-                    m.update(d.encode('utf8'))
-                if not lsh.query(m):
-                    lsh.insert(f"key_{i}", m)
+        elif method == 'Semántica (MinHash)':
+            lsh = MinHashLSH(threshold=threshold, num_perm=num_perm)
+            for i, example in enumerate(dataset):
+                text = example.get(text_col, "")
+                if text and isinstance(text, str) and text.strip():
+                    m = MinHash(num_perm=num_perm)
+                    for d in text.split():
+                        m.update(d.encode('utf8'))
+                    if not lsh.query(m):
+                        lsh.insert(f"key_{i}", m)
+                        yield example
+                else:
                     yield example
-            else:
-                yield example
+        else:
+             yield from dataset
+
+    new_ds = IterableDataset.from_generator(gen)
+    if hasattr(dataset, 'info'):
+        new_ds.info = dataset.info
+    elif hasattr(dataset, '_info'):
+        new_ds.info = dataset._info
+    return new_ds
 
 @spaces.GPU()
 def hf_login(token):
@@ -558,6 +547,10 @@ def _create_training_args(output_dir, repo_id, **kwargs):
         "adam_epsilon": float(kwargs.get('adam_epsilon', 1e-8)),
         "no_cuda": device == 'cpu'
     }
+
+    if kwargs.get('early_stopping_patience', 0) > 0 and kwargs.get('run_evaluation', False):
+        args_dict['early_stopping_patience'] = int(kwargs['early_stopping_patience'])
+        args_dict['load_best_model_at_end'] = True
     
     is_diffusion_task = kwargs.get('training_mode', '') in ["Text-to-Image (LoRA)", "DreamBooth LoRA (Text-to-Image)"]
     if is_diffusion_task:
@@ -652,35 +645,43 @@ def _find_all_linear_names(model, quantization_type):
     
     return list(lora_module_names.intersection(common_targets)) or list(lora_module_names)
 
-@spaces.GPU()
-def _conversation_formatting_func(example, tokenizer, **kwargs):
-    conv_col = ""
-    for key in ["messages", "conversations", "turns"]:
-        if key in example: conv_col = key; break
-    if not conv_col: return ""
-    conversation = example[conv_col]
-    if isinstance(conversation, str):
-        try: conversation = ast.literal_eval(conversation)
-        except: return ""
-    return tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
-
 @spaces.GPU()
 def _sft_formatting_func(example, text_col, tokenizer, **kwargs):
-    if kwargs.get('enable_cot_input') or kwargs.get('enable_tool_use_input'):
+    if kwargs.get('sft_format_style') == "Conversacional":
+        conv_col = ""
+        for key in ["messages", "conversations", "turns"]:
+            if key in example: conv_col = key; break
+        if not conv_col: return ""
+        conversation = example[conv_col]
+        if isinstance(conversation, str):
+            try: conversation = ast.literal_eval(conversation)
+            except: return ""
+        return tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
+
+    if kwargs.get('sft_format_style') == "Razonamiento/Herramientas":
         messages = []
         prompt = example.get(kwargs.get('prompt_col_input', 'prompt'), "")
         if prompt: messages.append({"role": "user", "content": prompt})
+        
         response_parts = []
-        if kwargs.get('enable_cot_input') and example.get(kwargs.get('reasoning_col_input', 'reasoning')): response_parts.append(f"<thinking>{example[kwargs.get('reasoning_col_input', 'reasoning')]}</thinking>")
-        if kwargs.get('enable_tool_use_input') and example.get(kwargs.get('tool_use_col_input', 'tools')): response_parts.append(f"<tool_code>{example[kwargs.get('tool_use_col_input', 'tools')]}</tool_code>")
-        if example.get(kwargs.get('response_col_input', 'response')): response_parts.append(example[kwargs.get('response_col_input', 'response')])
-        if response_parts: messages.append({"role": "assistant", "content": "\n".join(response_parts)})
+        if kwargs.get('enable_cot_input') and example.get(kwargs.get('reasoning_col_input', 'reasoning')):
+            response_parts.append(f"<thinking>{example[kwargs.get('reasoning_col_input', 'reasoning')]}</thinking>")
+        if kwargs.get('enable_tool_use_input') and example.get(kwargs.get('tool_use_col_input', 'tools')):
+            response_parts.append(f"<tool_code>{example[kwargs.get('tool_use_col_input', 'tools')]}</tool_code>")
+        if example.get(kwargs.get('response_col_input', 'response')):
+            response_parts.append(example[kwargs.get('response_col_input', 'response')])
+        
+        if response_parts:
+            messages.append({"role": "assistant", "content": "\n".join(response_parts)})
+        
         if messages:
             try:
                 return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
             except Exception as e:
-                logger.error(f"Error applying chat template: {e}.")
+                logger.error(f"Error aplicando la plantilla de chat: {e}.")
                 return "\n".join([m['content'] for m in messages])
+        return ""
+        
     return example.get(text_col, "")
 
 @spaces.GPU()
@@ -811,17 +812,20 @@ def train_sft_dpo(model_name, train_dataset, repo_id, update_logs_fn, model_card
                     eval_dataset = update
 
         TrainerClass = DPOTrainer if is_dpo else (DebiasingSFTTrainer if kwargs.get('enable_loss_reweighting') else SFTTrainer)
-        trainer_kwargs = {"model": model, "args": training_args, "train_dataset": train_dataset, "eval_dataset": eval_dataset, "peft_config": peft_config}
+        trainer_kwargs = {"model": model, "args": training_args, "train_dataset": train_dataset, "eval_dataset": eval_dataset, "tokenizer": tokenizer, "peft_config": peft_config}
 
         if is_dpo:
             trainer_kwargs.update({"beta": 0.1, "max_length": int(kwargs.get('block_size')), "max_prompt_length": int(kwargs.get('block_size')) // 2})
+            if train_dataset:
+                train_dataset = train_dataset.map(lambda ex: _dpo_formatting_func(ex, **kwargs))
             if eval_dataset:
                 eval_dataset = eval_dataset.map(lambda ex: _dpo_formatting_func(ex, **kwargs))
+            trainer_kwargs.update({"train_dataset": train_dataset, "eval_dataset": eval_dataset})
         else:
             sft_kwargs = kwargs.copy()
-            trainer_kwargs.update({"formatting_func": lambda ex, tc=text_col, skw=sft_kwargs: _sft_formatting_func(example=ex, tokenizer=tokenizer, text_col=tc, **skw)})
+            trainer_kwargs.update({"formatting_func": lambda ex: _sft_formatting_func(example=ex, tokenizer=tokenizer, text_col=text_col, **sft_kwargs), "max_seq_length": int(kwargs.get('block_size'))})
             if kwargs.get('enable_loss_reweighting'):
-                trainer_kwargs.update({'reweighting_terms': kwargs.get('reweighting_terms', '').split(','), 'reweighting_factor': kwargs.get('reweighting_factor', 2.0)})
+                trainer_kwargs.update({'reweighting_terms': kwargs.get('reweighting_terms', '').split(','), 'reweighting_factor': float(kwargs.get('reweighting_factor', 2.0))})
                 
         trainer = TrainerClass(**trainer_kwargs)
         final_model_path, final_metrics = yield from _run_trainer_and_upload(trainer, tokenizer, repo_id, update_logs_fn, model_card_content, **kwargs)
@@ -1277,7 +1281,7 @@ def _get_data_processing_pipeline(**kwargs):
     if kwargs.get('uploads'):
         uploaded_data_map = _load_uploaded_stream(kwargs.get('uploads'))
         if uploaded_data_map and uploaded_data_map["train"]:
-            train_dataset = Dataset.from_list(uploaded_data_map["train"])
+            train_dataset = IterableDataset.from_generator(lambda: iter(uploaded_data_map["train"]))
             uploaded_val_data = uploaded_data_map["validation"]
     
     if hf_ids:
@@ -1325,7 +1329,7 @@ def _get_data_processing_pipeline(**kwargs):
             
         dedup_method = kwargs.get('deduplication_method')
         if dedup_method != 'Ninguna':
-            train_dataset = DeduplicatedIterableDataset(
+            train_dataset = _create_deduplicated_iterable_dataset(
                 dataset=train_dataset,
                 text_col=text_col,
                 method=dedup_method,
@@ -1659,8 +1663,10 @@ def gradio_preview_data_wrapper(*args):
             formatted_text = ""
             if kwargs['training_mode'] == "DPO (Direct Preference Optimization)":
                 formatted_text = json.dumps(_dpo_formatting_func(example, **kwargs), indent=2, ensure_ascii=False)
-            else:
+            elif kwargs['training_mode'] == "Causal Language Modeling (SFT/LoRA)":
                 formatted_text = _sft_formatting_func(example, text_col, tokenizer, **kwargs)
+            else:
+                formatted_text = str(example)
 
             preview_samples.append(f"--- MUESTRA {i+1} ---\n{formatted_text}\n")
 
@@ -1696,14 +1702,19 @@ def toggle_task_specific_ui(training_mode):
         gr.update(visible=is_classification or is_ner),
         gr.update(visible=is_dpo),
         gr.update(visible=is_sft),
-        gr.update(visible=training_mode == "Text-to-Image (LoRA)"),
+        gr.update(visible=is_diffusion),
         gr.update(visible=training_mode == "DreamBooth LoRA (Text-to-Image)"),
         gr.update(visible=not is_diffusion),
         gr.update(visible=is_diffusion),
+        gr.update(visible=is_streaming),
         gr.update(visible=not is_streaming),
-        gr.update(visible=is_streaming)
     )
 
+@spaces.GPU()
+def toggle_sft_format_ui(format_style):
+    is_tool = format_style == "Razonamiento/Herramientas"
+    return gr.update(visible=is_tool)
+
 @spaces.GPU()
 def toggle_auto_modules_ui(is_auto):
     return gr.update(visible=not is_auto)
@@ -1784,10 +1795,10 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                         gradient_accumulation = gr.Textbox(label="Acumulación de Gradiente", value="8")
                     with gr.Row():
                         block_size = gr.Textbox(label="Longitud de Secuencia", value="1024")
+                        with gr.Group(visible=True) as max_steps_ui:
+                             max_steps = gr.Textbox(label="Máximos Pasos de Entrenamiento", value="100")
                         with gr.Group(visible=False) as epochs_ui:
                             epochs = gr.Textbox(label="Épocas", value="1")
-                        with gr.Group(visible=True) as max_steps_ui:
-                            max_steps = gr.Textbox(label="Máximos Pasos de Entrenamiento", value="100")
                     with gr.Row():
                         optimizer = gr.Dropdown(["adamw_torch", "adafactor", "sgd", "adagrad"], label="Optimizador", value="adamw_torch")
                         scheduler = gr.Dropdown(["cosine", "linear", "constant"], label="Planificador LR", value="cosine")
@@ -1799,6 +1810,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                          logging_steps = gr.Textbox(label="Pasos de Registro", value="10")
                          save_steps = gr.Textbox(label="Pasos de Guardado", value="50")
                          save_total_limit = gr.Textbox(label="Límite Total de Guardado", value="1")
+                         early_stopping_patience = gr.Number(label="Paciencia para Early Stopping (0 para desactivar)", value=0)
                          resume_from_checkpoint = gr.Checkbox(label="Reanudar desde Checkpoint", value=False)
                          with gr.Row():
                             adam_beta1 = gr.Textbox(label="Adam Beta1", value="0.9")
@@ -1851,25 +1863,34 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                         num_synthetic_samples = gr.Number(label="Número de Muestras", value=1000)
                 
                 with gr.Accordion("📝 Configuración de Formato y Tarea", open=False):
+                    with gr.Group(visible=True) as sft_ui:
+                        sft_format_style = gr.Radio(["Columna de Texto", "Conversacional", "Razonamiento/Herramientas"], label="Formato de Datos SFT", value="Columna de Texto")
+                        chat_template_jinja = gr.Textbox(label="Plantilla de Chat Jinja2 (opcional)", lines=5)
+                        with gr.Group(visible=False) as sft_tool_ui:
+                            enable_cot_input = gr.Checkbox(label="Habilitar Razonamiento (CoT)", value=True)
+                            enable_tool_use_input = gr.Checkbox(label="Habilitar Uso de Herramientas", value=True)
+                            prompt_col_input = gr.Textbox(label="Columna de Prompt/Usuario", value="prompt")
+                            response_col_input = gr.Textbox(label="Columna de Respuesta Final", value="response")
+                            reasoning_col_input = gr.Textbox(label="Columna de Razonamiento", value="reasoning")
+                            tool_use_col_input = gr.Textbox(label="Columna de Uso de Herramientas", value="tools")
+                    with gr.Group(visible=False) as dpo_ui:
+                        dpo_prompt_col_input = gr.Textbox(label="Columna de Prompt", value="prompt")
+                        dpo_chosen_col_input = gr.Textbox(label="Columna Elegida", value="chosen")
+                        dpo_rejected_col_input = gr.Textbox(label="Columna Rechazada", value="rejected")
+                    with gr.Group(visible=False) as classification_labels_ui:
+                        classification_labels = gr.Textbox(label="Etiquetas de Clasificación (csv)", placeholder="p.ej. positivo,negativo")
                     with gr.Group(visible=False) as diffusion_ui:
                         diffusion_resolution = gr.Slider(256, 1024, 512, step=64, label="Resolución")
                     with gr.Group(visible=False) as dreambooth_ui:
                         dreambooth_instance_prompt = gr.Textbox(label="Prompt de Instancia", placeholder="p.ej. 'foto de perro sks'")
                         dreambooth_train_text_encoder = gr.Checkbox(label="Entrenar Text Encoder", value=True)
-                    with gr.Group(visible=False) as classification_labels_ui:
-                        classification_labels = gr.Textbox(label="Etiquetas de Clasificación (csv)", placeholder="p.ej. positivo,negativo")
-                    with gr.Group(visible=False) as dpo_ui:
-                        prompt_col_input = gr.Textbox(label="Columna de Prompt", value="prompt")
-                        dpo_chosen_col_input = gr.Textbox(label="Columna Elegida", value="chosen")
-                        dpo_rejected_col_input = gr.Textbox(label="Columna Rechazada", value="rejected")
-                    with gr.Group(visible=True) as sft_ui:
-                        chat_template_jinja = gr.Textbox(label="Plantilla de Chat Jinja2 (opcional)", lines=5)
                 
                 with gr.Accordion("📊 Evaluación y Mitigación de Sesgos", open=False):
                     run_evaluation = gr.Checkbox(label="Ejecutar Evaluación", value=False)
                     run_perplexity_evaluation = gr.Checkbox(label="Calcular Perplejidad", value=True)
                     enable_loss_reweighting = gr.Checkbox(label="Habilitar Re-ponderación de Pérdida", value=False)
                     reweighting_terms = gr.Textbox(label="Términos para Re-ponderar (csv)", placeholder="sesgo,injusto")
+                    reweighting_factor = gr.Slider(1.1, 10.0, 2.0, label="Factor de Re-ponderación")
                     enable_cda = gr.Checkbox(label="Habilitar Aumentación Contrafactual (CDA)", value=False)
                     cda_json_config = gr.Textbox(label="Configuración CDA (JSON)", placeholder='[["ella", "él"], ["mujer", "hombre"]]')
 
@@ -1901,6 +1922,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             "adam_beta1": adam_beta1, "adam_beta2": adam_beta2, "adam_epsilon": adam_epsilon, 
             "disable_gradient_checkpointing": disable_gradient_checkpointing, "group_by_length": group_by_length, 
             "neftune_noise_alpha": neftune_noise_alpha, "optim_args": optim_args, "attn_implementation": attn_implementation, 
+            "early_stopping_patience": early_stopping_patience,
             "peft": peft, "quantization": quantization, "lora_r": lora_r, "lora_alpha": lora_alpha, 
             "lora_dropout": lora_dropout, "auto_find_target_modules": auto_find_target_modules, "target_modules": target_modules, 
             "modules_to_save": modules_to_save, "use_dora": use_dora, "use_rslora": use_rslora, "init_lora_weights": init_lora_weights, 
@@ -1912,10 +1934,14 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             "enable_back_translation": enable_back_translation, "bt_model_id": bt_model_id, 
             "bt_reverse_model_id": bt_reverse_model_id, "enable_synthetic_data": enable_synthetic_data, 
             "synthetic_model_id": synthetic_model_id, "num_synthetic_samples": num_synthetic_samples, 
-            "chat_template_jinja": chat_template_jinja, "prompt_col_input": prompt_col_input, "dpo_chosen_col_input": dpo_chosen_col_input, 
+            "sft_format_style": sft_format_style, "chat_template_jinja": chat_template_jinja, 
+            "enable_cot_input": enable_cot_input, "enable_tool_use_input": enable_tool_use_input,
+            "prompt_col_input": prompt_col_input, "response_col_input": response_col_input,
+            "reasoning_col_input": reasoning_col_input, "tool_use_col_input": tool_use_col_input,
+            "dpo_prompt_col_input": dpo_prompt_col_input, "dpo_chosen_col_input": dpo_chosen_col_input, 
             "dpo_rejected_col_input": dpo_rejected_col_input, "classification_labels": classification_labels, 
             "diffusion_resolution": diffusion_resolution, "run_evaluation": run_evaluation, "run_perplexity_evaluation": run_perplexity_evaluation, 
-            "enable_loss_reweighting": enable_loss_reweighting, "reweighting_terms": reweighting_terms, 
+            "enable_loss_reweighting": enable_loss_reweighting, "reweighting_terms": reweighting_terms, "reweighting_factor": reweighting_factor,
             "wandb_api_key_input": wandb_api_key_input, "wandb_project_input": wandb_project_input,
             "dreambooth_instance_prompt": dreambooth_instance_prompt,
             "dreambooth_train_text_encoder": dreambooth_train_text_encoder
@@ -1939,7 +1965,13 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
         training_mode.change(
             toggle_task_specific_ui,
             inputs=[training_mode],
-            outputs=[classification_labels_ui, dpo_ui, sft_ui, diffusion_ui, dreambooth_ui, peft_accordion, epochs_ui, max_steps_ui]
+            outputs=[classification_labels_ui, dpo_ui, sft_ui, diffusion_ui, dreambooth_ui, peft_accordion, epochs_ui, max_steps_ui, peft_accordion]
+        )
+        
+        sft_format_style.change(
+            toggle_sft_format_ui,
+            inputs=[sft_format_style],
+            outputs=[sft_tool_ui]
         )
 
         auto_find_target_modules.change(