Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 os.system("pip install -U transformers peft accelerate trl bitsandbytes datasets diffusers")
 os.system("pip install spaces-0.1.0-py3-none-any.whl")
-import os
 import io
 import json
 import tempfile
@@ -17,9 +16,6 @@ import re
 import ast
 from itertools import islice
 from pathlib import Path
-from collections import defaultdict
-from datetime import datetime
-
 import torch
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
@@ -33,14 +29,15 @@ from langdetect import detect_langs
 import textstat
 from datasketch import MinHash, MinHashLSH
 import gradio as gr
-from datasets import load_dataset, IterableDataset, Dataset, DatasetDict, interleave_datasets, Audio
+import spaces
+from datasets import load_dataset, IterableDataset, Dataset, DatasetDict
 from huggingface_hub import login, whoami, create_repo, upload_folder, HfApi
 from transformers import (
     AutoModelForCausalLM, AutoTokenizer, AutoConfig, TrainingArguments, Trainer,
     AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer,
-    SpeechT5ForTextToSpeech, SpeechT5Processor, SpeechT5HifiGan, AutoModelForImageClassification,
+    AutoModelForImageClassification,
     AutoImageProcessor, AutoModelForAudioClassification, AutoFeatureExtractor, AutoModelForTokenClassification,
-    DataCollatorForTokenClassification, AutoModelForQuestionAnswering, AutoModelForSpeechSeq2Seq,
+    DataCollatorForTokenClassification, AutoModelForQuestionAnswering,
     AutoProcessor, DataCollatorWithPadding, pipeline, CLIPTextModel, CLIPTokenizer,
     DataCollatorForSeq2Seq, AutoModelForSequenceClassification, BitsAndBytesConfig,
     LlamaConfig, LlamaForCausalLM, MistralConfig, MistralForCausalLM, GemmaConfig, GemmaForCausalLM, GPT2Config, GPT2LMHeadModel,
@@ -57,8 +54,8 @@ from diffusers import (
 )
 import evaluate as hf_evaluate
 from jinja2 import Template
+from collections import defaultdict
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 if torch.cuda.is_available():
@@ -97,8 +94,7 @@ TASK_TO_PIPELINE_MAP = {
     "DreamBooth LoRA (Text-to-Image)": "text-to-image",
 }
 
-MODEL_CARD_TEMPLATE = """
----
+MODEL_CARD_TEMPLATE = """---
 language: es
 license: apache-2.0
 tags:
@@ -136,8 +132,7 @@ Este modelo es una versión afinada de [{base_model}](https://huggingface.co/{ba
 - Gradio
 """
 
-DATASET_CARD_TEMPLATE = """
----
+DATASET_CARD_TEMPLATE = """---
 license: mit
 ---
 
@@ -169,52 +164,6 @@ class DebiasingSFTTrainer(SFTTrainer):
                     break
         return (loss, outputs) if return_outputs else loss
 
-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)
-                    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)
-                    yield example
-            else:
-                yield example
-
 def hf_login(token):
     if not token:
         return "Por favor, introduce un token."
@@ -330,6 +279,8 @@ def _load_hf_streaming(ids, split="train", probabilities=None):
     if probabilities and len(probabilities) != len(streams):
         logger.warning(f"Number of probabilities ({len(probabilities)}) does not match number of valid datasets ({len(streams)}). Ignoring weights.")
         probabilities = None
+    
+    from datasets import interleave_datasets
     return interleave_datasets(streams, probabilities=probabilities)
 
 def _load_uploaded_stream(files):
@@ -517,7 +468,6 @@ def _create_training_args(output_dir, repo_id, **kwargs):
         "save_strategy": "steps",
         "logging_steps": int(kwargs.get('logging_steps', 10)),
         "save_steps": int(kwargs.get('save_steps', 50)),
-        "evaluation_strategy": "steps" if kwargs.get('run_evaluation', False) else "no",
         "eval_steps": int(kwargs.get('save_steps', 50)) if kwargs.get('run_evaluation', False) else None,
         "learning_rate": float(kwargs.get('learning_rate', 2e-5)),
         "fp16": kwargs.get('mixed_precision') == 'fp16' and device == 'cuda',
@@ -727,7 +677,6 @@ def _run_trainer_and_upload(trainer, tokenizer, repo_id, update_logs_fn, model_c
         eval_logs = [log for log in trainer.state.log_history if 'eval_loss' in log]
         if eval_logs:
             final_metrics = eval_logs[-1]
-            final_metrics = {k.replace('eval_', ''): v for k, v in final_metrics.items()}
 
     yield update_logs_fn("Entrenamiento finalizado.", "Guardando")
     output_dir = trainer.args.output_dir
@@ -777,16 +726,16 @@ def train_sft_dpo(model_name, train_dataset, repo_id, update_logs_fn, model_card
         if kwargs.get('run_evaluation'):
             eval_dataset_gen = _get_eval_dataset(kwargs.get('datasets_hf_text').split(","), kwargs.get('eval_dataset_hf'), kwargs.get('uploaded_val_data'), update_logs_fn)
             for update in eval_dataset_gen:
-                if isinstance(update, dict):
+                if isinstance(update, tuple):
                     yield update
                 else:
                     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, "tokenizer": tokenizer, "max_seq_length": int(kwargs.get('block_size'))}
+        trainer_kwargs = {"model": model, "args": training_args, "train_dataset": train_dataset, "eval_dataset": eval_dataset, "peft_config": peft_config, "tokenizer": tokenizer}
 
         if is_dpo:
-            trainer_kwargs.update({"beta": 0.1, "max_prompt_length": int(kwargs.get('block_size')) // 2})
+            trainer_kwargs.update({"beta": 0.1, "max_length": int(kwargs.get('block_size')), "max_prompt_length": int(kwargs.get('block_size')) // 2})
             if eval_dataset:
                 eval_dataset = eval_dataset.map(lambda ex: _dpo_formatting_func(ex, **kwargs))
         else:
@@ -794,8 +743,18 @@ def train_sft_dpo(model_name, train_dataset, repo_id, update_logs_fn, model_card
             trainer_kwargs.update({"formatting_func": lambda ex: _sft_formatting_func(example=ex, tokenizer=tokenizer, text_col=text_col, **sft_kwargs)})
             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 = TrainerClass(**trainer_kwargs)
+        
+        try:
+            trainer = TrainerClass(**trainer_kwargs)
+        except TypeError as e:
+            if "unexpected keyword argument 'tokenizer'" in str(e):
+                logger.warning("Caught TypeError for tokenizer argument. Retrying without it for TRL compatibility.")
+                trainer_kwargs.pop("tokenizer", None)
+                trainer = TrainerClass(**trainer_kwargs)
+                trainer.tokenizer = tokenizer
+            else:
+                raise e
+        
         final_model_path, final_metrics = yield from _run_trainer_and_upload(trainer, tokenizer, repo_id, update_logs_fn, model_card_content, **kwargs)
         return final_model_path, final_metrics
 
@@ -812,26 +771,23 @@ def train_sequence_classification(model_name, train_dataset, repo_id, update_log
         tokenizer_id = kwargs.get('tokenizer_name') or model_name
         yield update_logs_fn(f"Cargando tokenizer '{tokenizer_id}'...", "Configuración")
         tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
         
         yield update_logs_fn(f"Cargando modelo '{model_name}'...", "Configuración")
         model = _generic_model_loader(model_name, AutoModelForSequenceClassification, num_labels=len(labels), label2id=label2id, id2label=id2label, **kwargs)
-        model.config.pad_token_id = tokenizer.pad_token_id
 
         def preprocess(examples):
             return tokenizer(examples[kwargs['text_col']], truncation=True, max_length=512)
-        train_dataset = train_dataset.map(preprocess, batched=True)
+        train_dataset = train_dataset.map(preprocess)
 
         eval_dataset = None
         if kwargs.get('run_evaluation'):
             eval_dataset_gen = _get_eval_dataset(kwargs.get('datasets_hf_text').split(","), kwargs.get('eval_dataset_hf'), kwargs.get('uploaded_val_data'), update_logs_fn)
             for update in eval_dataset_gen:
-                if isinstance(update, dict):
+                if isinstance(update, tuple):
                     yield update
                 else:
                     eval_dataset = update
-            if eval_dataset: eval_dataset = eval_dataset.map(preprocess, batched=True)
+            if eval_dataset: eval_dataset = eval_dataset.map(preprocess)
 
         metric = hf_evaluate.load("accuracy")
         def compute_metrics(eval_pred):
@@ -887,7 +843,7 @@ def train_token_classification(model_name, train_dataset, repo_id, update_logs_f
         if kwargs.get('run_evaluation'):
             eval_dataset_gen = _get_eval_dataset(kwargs.get('datasets_hf_text').split(","), kwargs.get('eval_dataset_hf'), kwargs.get('uploaded_val_data'), update_logs_fn)
             for update in eval_dataset_gen:
-                if isinstance(update, dict):
+                if isinstance(update, tuple):
                     yield update
                 else:
                     eval_dataset = update
@@ -979,7 +935,7 @@ def train_question_answering(model_name, train_dataset, repo_id, update_logs_fn,
             eval_dataset_raw_gen = _get_eval_dataset(kwargs.get('datasets_hf_text').split(","), kwargs.get('eval_dataset_hf'), kwargs.get('uploaded_val_data'), update_logs_fn)
             eval_dataset_raw = None
             for update in eval_dataset_raw_gen:
-                if isinstance(update, dict):
+                if isinstance(update, tuple):
                     yield update
                 else:
                     eval_dataset_raw = update
@@ -1023,7 +979,7 @@ def train_seq2seq(model_name, train_dataset, repo_id, update_logs_fn, model_card
         if kwargs.get('run_evaluation'):
             eval_dataset_gen = _get_eval_dataset(kwargs.get('datasets_hf_text').split(","), kwargs.get('eval_dataset_hf'), kwargs.get('uploaded_val_data'), update_logs_fn)
             for update in eval_dataset_gen:
-                if isinstance(update, dict):
+                if isinstance(update, tuple):
                     yield update
                 else:
                     eval_dataset = update
@@ -1067,30 +1023,22 @@ def train_text_to_image(model_name, train_dataset, repo_id, update_logs_fn, mode
 
     yield update_logs_fn("Configurando componentes de Diffusers...", "Text-to-Image (LoRA)")
     tokenizer = CLIPTokenizer.from_pretrained(model_name, subfolder="tokenizer")
-    text_encoder = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder", torch_dtype=torch_dtype_auto)
-    vae = AutoencoderKL.from_pretrained(model_name, subfolder="vae", torch_dtype=torch_dtype_auto)
-    unet = UNet2DConditionModel.from_pretrained(model_name, subfolder="unet", torch_dtype=torch_dtype_auto)
+    text_encoder = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder")
+    vae = AutoencoderKL.from_pretrained(model_name, subfolder="vae")
+    unet = UNet2DConditionModel.from_pretrained(model_name, subfolder="unet")
     noise_scheduler = DDPMScheduler.from_pretrained(model_name, subfolder="scheduler")
     
     vae.requires_grad_(False)
     text_encoder.requires_grad_(False)
     unet.train()
 
-    yield update_logs_fn("Agregando adaptadores LoRA al UNet...", "Text-to-Image (LoRA)")
+    yield update_logs_fn("Agregando adaptadores LoRA al modelo...", "Text-to-Image (LoRA)")
     unet_lora_config = LoraConfig(
         r=int(kwargs.get('lora_r', 16)), lora_alpha=int(kwargs.get('lora_alpha', 32)),
         target_modules=["to_q", "to_k", "to_v", "to_out.0"],
     )
     unet.add_adapter(unet_lora_config)
 
-    if kwargs.get('dreambooth_train_text_encoder', False):
-        yield update_logs_fn("Agregando adaptadores LoRA al Text Encoder...", "DreamBooth LoRA")
-        text_encoder_lora_config = LoraConfig(
-            r=int(kwargs.get('lora_r', 16)), lora_alpha=int(kwargs.get('lora_alpha', 32)),
-            target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
-        )
-        text_encoder.add_adapter(text_encoder_lora_config)
-
     yield update_logs_fn("Procesando dataset de imágenes...", "Text-to-Image (LoRA)")
     resolution = int(kwargs.get('diffusion_resolution', 512))
     
@@ -1102,7 +1050,7 @@ def train_text_to_image(model_name, train_dataset, repo_id, update_logs_fn, mode
     ])
 
     def preprocess_train(examples):
-        images = [image.convert("RGB") for image in examples[kwargs.get('image_col', 'image')]]
+        images = [Image.open(image).convert("RGB") for image in examples[kwargs.get('image_col', 'image')]]
         examples["pixel_values"] = [train_transforms(image) for image in images]
         examples["input_ids"] = tokenizer(examples[kwargs.get('text_col', 'text')], max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt").input_ids
         return examples
@@ -1116,17 +1064,14 @@ def train_text_to_image(model_name, train_dataset, repo_id, update_logs_fn, mode
 
     def collate_fn(examples):
         pixel_values = torch.stack([example["pixel_values"] for example in examples])
-        input_ids = torch.stack([e["input_ids"][0] for e in examples])
+        input_ids = torch.stack([example["input_ids"] for example in examples])
         return {"pixel_values": pixel_values, "input_ids": input_ids}
         
     train_dataloader = DataLoader(processed_dataset, shuffle=True, collate_fn=collate_fn, batch_size=int(kwargs.get('batch_size', 1)))
-    
-    params_to_optimize = list(unet.parameters())
-    if kwargs.get('dreambooth_train_text_encoder', False):
-        params_to_optimize += list(text_encoder.parameters())
 
+    yield update_logs_fn("Configurando optimizador y planificador...", "Text-to-Image (LoRA)")
     optimizer = torch.optim.AdamW(
-        params_to_optimize, lr=float(kwargs.get('learning_rate', 2e-5)),
+        unet.parameters(), lr=float(kwargs.get('learning_rate', 2e-5)),
         betas=(float(kwargs.get('adam_beta1', 0.9)), float(kwargs.get('adam_beta2', 0.999))),
         weight_decay=float(kwargs.get('weight_decay', 0.01)),
         eps=float(kwargs.get('adam_epsilon', 1e-8)),
@@ -1142,34 +1087,36 @@ def train_text_to_image(model_name, train_dataset, repo_id, update_logs_fn, mode
         num_training_steps=max_train_steps,
     )
     
-    unet, text_encoder, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
-        unet, text_encoder, optimizer, train_dataloader, lr_scheduler
+    unet, optimizer, train_dataloader, lr_scheduler, text_encoder, vae = accelerator.prepare(
+        unet, optimizer, train_dataloader, lr_scheduler, text_encoder, vae
     )
     
+    text_encoder.to(accelerator.device, dtype=torch_dtype_auto)
     vae.to(accelerator.device, dtype=torch_dtype_auto)
 
+    yield update_logs_fn("Iniciando bucle de entrenamiento de difusión...", "Text-to-Image (LoRA)")
     global_step = 0
     final_loss = 0
     for epoch in range(num_epochs):
         for step, batch in enumerate(train_dataloader):
             with accelerator.accumulate(unet):
-                latents = vae.encode(batch["pixel_values"].to(dtype=torch_dtype_auto)).latent_dist.sample()
+                latents = vae.encode(batch["pixel_values"].to(torch_dtype_auto)).latent_dist.sample()
                 latents = latents * vae.config.scaling_factor
                 noise = torch.randn_like(latents)
                 bsz = latents.shape[0]
-                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device).long()
+                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
+                timesteps = timesteps.long()
+
                 noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
-                encoder_hidden_states = text_encoder(batch["input_ids"])[0]
+                encoder_hidden_states = text_encoder(batch["input_ids"].to(accelerator.device))[0]
+
                 noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
                 loss = F.mse_loss(noise_pred.float(), noise.float(), reduction="mean")
                 final_loss = loss.detach().item()
                 
                 accelerator.backward(loss)
                 if accelerator.sync_gradients:
-                    params_to_clip = list(unet.parameters())
-                    if kwargs.get('dreambooth_train_text_encoder', False):
-                        params_to_clip += list(text_encoder.parameters())
-                    accelerator.clip_grad_norm_(params_to_clip, float(kwargs.get('max_grad_norm', 1.0)))
+                    accelerator.clip_grad_norm_(unet.parameters(), float(kwargs.get('max_grad_norm', 1.0)))
                 
                 optimizer.step()
                 lr_scheduler.step()
@@ -1177,21 +1124,16 @@ def train_text_to_image(model_name, train_dataset, repo_id, update_logs_fn, mode
 
             if accelerator.is_main_process:
                 if global_step % int(kwargs.get('logging_steps', 10)) == 0:
-                    yield update_logs_fn(f"Epoch {epoch}, Step {step}, Loss: {final_loss:.4f}", "Entrenando Difusión")
+                    yield update_logs_fn(f"Epoch {epoch}, Step {step}, Loss: {final_loss}", "Text-to-Image (LoRA)")
             global_step += 1
-            if global_step >= max_train_steps:
-                break
-        if global_step >= max_train_steps:
-            break
 
+    yield update_logs_fn("Entrenamiento completado, guardando modelo...", "Text-to-Image (LoRA)")
     accelerator.wait_for_everyone()
     if accelerator.is_main_process:
-        pipeline = StableDiffusionText2ImagePipeline.from_pretrained(
-            model_name,
-            unet=accelerator.unwrap_model(unet),
-            text_encoder=accelerator.unwrap_model(text_encoder),
-            torch_dtype=torch_dtype_auto,
-        )
+        unwrapped_unet = accelerator.unwrap_model(unet)
+        
+        pipeline = StableDiffusionText2ImagePipeline.from_pretrained(model_name, torch_dtype=torch_dtype_auto)
+        pipeline.unet.load_state_dict(unwrapped_unet.state_dict())
         pipeline.save_pretrained(output_dir)
         
         with open(os.path.join(output_dir, "README.md"), "w", encoding="utf-8") as f:
@@ -1206,6 +1148,7 @@ def train_text_to_image(model_name, train_dataset, repo_id, update_logs_fn, mode
         torch.cuda.empty_cache()
     return output_dir, {"final_loss": final_loss}
 
+
 def train_dreambooth_lora(model_name, train_dataset, repo_id, update_logs_fn, model_card_content, **kwargs):
     if device == 'cpu':
         raise ValueError("El entrenamiento de DreamBooth solo es compatible con GPU CUDA.")
@@ -1220,7 +1163,7 @@ def train_dreambooth_lora(model_name, train_dataset, repo_id, update_logs_fn, mo
 
     train_dataset = train_dataset.map(add_prompt)
     
-    yield update_logs_fn(f"Usando el prompt de instancia para todas las imágenes: '{dreambooth_prompt}'", "DreamBooth LoRA")
+    yield update_logs_fn(f"Usando el prompt de instancia para todas las imágenes: '{dreambooth_prompt}'", "DreamBooth LoRA (Text-to-Image)")
     
     final_model_path, final_metrics = yield from train_text_to_image(model_name, train_dataset, repo_id, update_logs_fn, model_card_content, **kwargs)
     return final_model_path, final_metrics
@@ -1251,8 +1194,9 @@ def _get_data_processing_pipeline(**kwargs):
             if train_dataset is None:
                 train_dataset = hf_train_dataset
             else:
+                from datasets import interleave_datasets
                 all_streams = [train_dataset, hf_train_dataset]
-                all_probs = [0.5, 0.5]
+                all_probs = [0.5, 0.5] if not probabilities else [probabilities] + probabilities[1:]
                 train_dataset = interleave_datasets(all_streams, probabilities=all_probs)
     
     if train_dataset is None:
@@ -1262,8 +1206,7 @@ def _get_data_processing_pipeline(**kwargs):
     text_col, image_col, audio_col, label_col = _guess_columns(first_example)
     kwargs.update({'text_col': text_col, 'image_col': image_col, 'audio_col': audio_col, 'label_col': label_col, 'uploaded_val_data': uploaded_val_data})
     
-    is_text_task = kwargs['training_mode'] not in ["DreamBooth LoRA (Text-to-Image)", "Text-to-Image (LoRA)", "Image Classification (Vision)", "Audio Classification (Speech)"]
-    if is_text_task:
+    if kwargs['training_mode'] not in ["DreamBooth LoRA (Text-to-Image)", "Text-to-Image (LoRA)"]:
         if any([kwargs.get('remove_html_tags'), kwargs.get('normalize_whitespace'), kwargs.get('remove_urls_emails'), kwargs.get('redact_pii')]):
             clean_kwargs = {k:v for k,v in kwargs.items() if k in ['remove_html_tags', 'normalize_whitespace', 'remove_urls_emails', 'redact_pii']}
             train_dataset = train_dataset.map(lambda ex: _clean_text(ex, text_col, **clean_kwargs))
@@ -1286,13 +1229,35 @@ def _get_data_processing_pipeline(**kwargs):
             
         dedup_method = kwargs.get('deduplication_method')
         if dedup_method != 'Ninguna':
-            train_dataset = DeduplicatedIterableDataset(
-                dataset=train_dataset,
-                text_col=text_col,
-                method=dedup_method,
-                threshold=kwargs.get('minhash_threshold', 0.85),
-                num_perm=int(kwargs.get('minhash_num_perm', 128))
-            )
+            base_iterator = train_dataset
+            if dedup_method == 'Exacta':
+                def dedup_generator_exact():
+                    seen_texts = set()
+                    for example in base_iterator:
+                        text = example.get(text_col, "")
+                        if not isinstance(text, str) or text not in seen_texts:
+                            if isinstance(text, str) and text:
+                                seen_texts.add(text)
+                            yield example
+                train_dataset = IterableDataset.from_generator(dedup_generator_exact)
+            elif dedup_method == 'Semántica (MinHash)':
+                threshold = kwargs.get('minhash_threshold', 0.85)
+                num_perm = int(kwargs.get('minhash_num_perm', 128))
+                def dedup_generator_minhash():
+                    lsh = MinHashLSH(threshold=threshold, num_perm=num_perm)
+                    for i, example in enumerate(base_iterator):
+                        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
+                train_dataset = IterableDataset.from_generator(dedup_generator_minhash)
+
 
     return train_dataset, kwargs
 
@@ -1385,6 +1350,7 @@ def _train_and_upload(**kwargs):
                  raise Exception(f"No se pudo cargar el tokenizer base '{tokenizer_id}' para el modelo desde cero: {e}")
             base_model_id_for_training = temp_model_dir
             kwargs["peft"] = False
+            kwargs["merge_adapter"] = False
             kwargs['tokenizer_name'] = temp_model_dir 
             yield update_logs(f"Modelo {architecture} inicializado en {temp_model_dir}.", "Modelo Cero") + (gr.update(), gr.update())
 
@@ -1397,6 +1363,7 @@ def _train_and_upload(**kwargs):
             os.environ["WANDB_PROJECT"] = kwargs.get('wandb_project_input') or f"{repo_base}"
             os.environ["WANDB_LOG_MODEL"] = "checkpoint"
         
+        from datetime import datetime
         model_card_content = MODEL_CARD_TEMPLATE.format(
             repo_id=repo_id, base_model=model_name, base_model_name=model_name.split('/')[-1],
             training_mode=kwargs.get('training_mode'),
@@ -1423,11 +1390,8 @@ def _train_and_upload(**kwargs):
             train_generator = train_func(base_model_id_for_training, train_dataset, repo_id, update_logs, model_card_content, **kwargs)
             while True:
                 try:
-                    update = next(train_generator)
-                    if isinstance(update, tuple) and len(update) == 4:
-                         yield update + (gr.update(), gr.update())
-                    else:
-                        pass
+                    update_tuple = next(train_generator)
+                    yield update_tuple + (gr.update(), gr.update())
                 except StopIteration as e:
                     final_model_path, final_metrics = e.value
                     break
@@ -1441,7 +1405,7 @@ def _train_and_upload(**kwargs):
             eval_dataset_perp = None
             eval_gen = _get_eval_dataset(kwargs.get('datasets_hf_text').split(","), kwargs.get('eval_dataset_hf'), kwargs.get('uploaded_val_data'), lambda m, p: update_logs(m, p))
             for update in eval_gen:
-                if isinstance(update, dict):
+                if isinstance(update, tuple):
                     yield update + (gr.update(), gr.update())
                 else:
                     eval_dataset_perp = update
@@ -1561,6 +1525,7 @@ def create_and_upload_dataset(hf_token, repo_name, creation_type, synth_model, s
                 for item in all_data:
                     f.write(json.dumps(item, ensure_ascii=False) + "\n")
 
+            from datetime import datetime
             readme_content = DATASET_CARD_TEMPLATE.format(
                 repo_id=repo_id,
                 creation_type=creation_type,
@@ -1598,13 +1563,9 @@ def gradio_preview_data_wrapper(*args):
         dataset, processed_kwargs = _get_data_processing_pipeline(**kwargs)
         text_col = processed_kwargs.get('text_col')
         
-        model_id_for_tokenizer = kwargs.get('model_base_input')
-        if not model_id_for_tokenizer:
-            raise ValueError("Se necesita un ID de modelo base para cargar el tokenizer para la vista previa.")
-            
-        tokenizer_id = kwargs.get('tokenizer_name') or model_id_for_tokenizer
         tokenizer = AutoTokenizer.from_pretrained(
-            tokenizer_id, trust_remote_code=True, use_fast=False
+            kwargs.get('tokenizer_name') or kwargs.get('model_base_input') or 'gpt2',
+            trust_remote_code=True, use_fast=False
         )
         if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token
         if kwargs.get('chat_template_jinja', '').strip(): tokenizer.chat_template = kwargs['chat_template_jinja']
@@ -1613,15 +1574,15 @@ def gradio_preview_data_wrapper(*args):
         for i, example in enumerate(islice(dataset, 5)):
             formatted_text = ""
             if kwargs['training_mode'] == "DPO (Direct Preference Optimization)":
-                formatted_text = json.dumps(_dpo_formatting_func(example, **kwargs), indent=2, ensure_ascii=False)
+                formatted_text = json.dumps(_dpo_formatting_func(example, **kwargs), indent=2)
             else:
                 formatted_text = _sft_formatting_func(example, text_col, tokenizer, **kwargs)
 
             preview_samples.append(f"--- MUESTRA {i+1} ---\n{formatted_text}\n")
 
         preview_text = "\n".join(preview_samples)
-        if not preview_samples:
-            preview_text = "No se pudieron generar muestras. Revisa la configuración del dataset, los filtros y el formato."
+        if not preview_text:
+            preview_text = "No se pudieron generar muestras. Revisa la configuración del dataset y el formato."
         yield preview_text
 
     except Exception as e:
@@ -1643,7 +1604,6 @@ def toggle_task_specific_ui(training_mode):
     is_sft = "Causal" in training_mode
     is_ner = "Token Classification" in training_mode
     is_diffusion = training_mode in ["Text-to-Image (LoRA)", "DreamBooth LoRA (Text-to-Image)"]
-    is_streaming = not is_diffusion
     
     return (
         gr.update(visible=is_classification or is_ner),
@@ -1653,10 +1613,10 @@ def toggle_task_specific_ui(training_mode):
         gr.update(visible=training_mode == "DreamBooth LoRA (Text-to-Image)"),
         gr.update(visible=not is_diffusion),
         gr.update(visible=is_diffusion),
-        gr.update(visible=not is_streaming),
-        gr.update(visible=is_streaming)
+        gr.update(visible=not is_diffusion),
     )
 
+
 def toggle_auto_modules_ui(is_auto):
     return gr.update(visible=not is_auto)
 
@@ -1690,7 +1650,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                     dset_file_uploads = gr.File(label="Subir Archivos (.jsonl, .csv, .txt)", file_count="multiple")
                 dset_create_button = gr.Button("Crear y Subir Dataset", variant="primary")
             with gr.Column(scale=2):
-                dset_status_output = gr.Textbox(label="Estado", lines=10, interactive=False)
+                dset_status_output = gr.Textbox(label="Estado", lines=10)
                 dset_link_output = gr.Markdown()
         
         dset_creation_type.change(toggle_dataset_creator_ui, inputs=[dset_creation_type], outputs=[dset_synth_group, dset_file_group])
@@ -1746,7 +1706,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                     with gr.Accordion("Avanzados", open=False):
                          warmup_ratio = gr.Slider(0.0, 0.5, 0.03, label="Ratio de Calentamiento")
                          weight_decay = gr.Textbox(label="Decaimiento de Peso", value="0.01")
-                         max_grad_norm = gr.Textbox(label="Norma Máxima de Gradiente", value="1.0")
+                         max_grad_norm = gr.Textbox(label="Norma Máxima de Gradiente", value="0.3")
                          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")
@@ -1806,6 +1766,9 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                         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_class_prompt = gr.Textbox(label="Prompt de Clase (Opcional)", placeholder="p.ej. 'foto de perro'")
+                        dreambooth_num_class_images = gr.Slider(0, 1000, 100, step=10, label="Nº de Imágenes de Clase")
+                        dreambooth_prior_loss_weight = gr.Slider(0.0, 2.0, 1.0, label="Peso de Pérdida a Priori")
                         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")
@@ -1824,6 +1787,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
                     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"]]')
 
+
                 with gr.Accordion("🔌 Integraciones", open=False):
                     wandb_api_key_input = gr.Textbox(label="Clave API de W&B", type="password")
                     wandb_project_input = gr.Textbox(label="Proyecto W&B")
@@ -1868,7 +1832,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             "diffusion_resolution": diffusion_resolution, "run_evaluation": run_evaluation, "run_perplexity_evaluation": run_perplexity_evaluation, 
             "enable_loss_reweighting": enable_loss_reweighting, "reweighting_terms": reweighting_terms, 
             "wandb_api_key_input": wandb_api_key_input, "wandb_project_input": wandb_project_input,
-            "dreambooth_instance_prompt": dreambooth_instance_prompt,
+            "dreambooth_instance_prompt": dreambooth_instance_prompt, "dreambooth_class_prompt": dreambooth_class_prompt,
+            "dreambooth_num_class_images": dreambooth_num_class_images, "dreambooth_prior_loss_weight": dreambooth_prior_loss_weight,
             "dreambooth_train_text_encoder": dreambooth_train_text_encoder
         }
         
@@ -1940,4 +1905,4 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             outputs=[inf_text_out, inf_model_id, inf_text_in, inf_context_in, inf_image_in, inf_audio_in]
         )
 
-demo.queue().launch(debug=True)
+demo.queue().launch(server_name="0.0.0.0", server_port=7860)