import os
os.system("pip install -U bitsandbytes diffusers torchaudio torchvision torch transformers peft accelerate trl datasets")
os.system("pip install spaces-0.1.0-py3-none-any.whl")
import io
import json
import tempfile
import string
import gc
import math
import uuid
import logging
import traceback
import importlib
import random
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 as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
import numpy as np
import pandas as pd
import accelerate
from PIL import Image
import torchvision
import torchvision.transforms as T
from torchvision import transforms
import torchaudio
from bs4 import BeautifulSoup
from langdetect import detect_langs
import textstat
from datasketch import MinHash, MinHashLSH
import gradio as gr
from datasets import load_dataset, IterableDataset, Dataset as HFDataset, DatasetDict, interleave_datasets, Audio
from huggingface_hub import login, whoami, create_repo, upload_folder, HfApi, hf_hub_download, list_repo_files
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, AutoConfig, TrainingArguments, Trainer,
    AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer,
    AutoModelForImageClassification,
    AutoImageProcessor, AutoModelForAudioClassification, AutoFeatureExtractor, AutoModelForTokenClassification,
    DataCollatorForTokenClassification, AutoModelForQuestionAnswering, AutoModelForSpeechSeq2Seq,
    AutoProcessor, DataCollatorWithPadding, pipeline,
    DataCollatorForSeq2Seq, AutoModelForSequenceClassification,
    LlamaConfig, LlamaForCausalLM, MistralConfig, MistralForCausalLM, GemmaConfig, GemmaForCausalLM, GPT2Config, GPT2LMHeadModel,
    PhiConfig, PhiForCausalLM, Qwen2Config, Qwen2ForCausalLM,
    DataCollatorForLanguageModeling, DefaultDataCollator, Adafactor
)
from peft import LoraConfig, get_peft_model, PeftModel
from trl import SFTTrainer, DPOTrainer
import evaluate as hf_evaluate
from jinja2 import Template
import spaces
from tqdm.auto import tqdm
from diffusers import (
    UNet2DConditionModel, DDPMScheduler, AutoencoderKL,
    get_scheduler as get_diffusers_scheduler, StableDiffusionPipeline as StableDiffusionText2ImagePipeline,
    StableDiffusionImg2ImgPipeline as StableDiffusionImage2ImagePipeline
)

logger = logging.getLogger(__name__)
torch_dtype_auto = torch.float32

def _sanitize_model_name_for_yaml(model_name):
    """Sanitize model name to be YAML-safe for use in tags"""
    # Extract the last part after '/' if present
    name = model_name.split('/')[-1] if '/' in model_name else model_name
    # Remove special characters that might break YAML
    # Keep only alphanumeric, hyphens, underscores, and dots
    sanitized = re.sub(r'[^a-zA-Z0-9\-_\.]', '-', name)
    # Ensure it's not empty
    return sanitized if sanitized else "model"

ARCHITECTURE_MAP = {"Llama": (LlamaConfig, LlamaForCausalLM), "Mistral": (MistralConfig, MistralForCausalLM), "Gemma": (GemmaConfig, GemmaForCausalLM), "GPT2": (GPT2Config, GPT2LMHeadModel), "Phi": (PhiConfig, PhiForCausalLM), "Qwen2": (Qwen2Config, Qwen2ForCausalLM)}
SCRATCH_TOKENIZER_MAP = {"Llama": "meta-llama/Llama-2-7b-hf", "Mistral": "mistralai/Mistral-7B-v0.1", "Gemma": "google/gemma-2b", "GPT2": "gpt2", "Phi": "microsoft/phi-2", "Qwen2": "Qwen/Qwen2-0.5B"}
TRAINING_MODES = [
    "Causal Language Modeling (SFT/LoRA)",
    "DPO (Direct Preference Optimization)",
    "Question Answering (Text)",
    "Token Classification (NER)",
    "Sequence Classification (Text)",
    "Text-to-Image Generation",
    "Image Classification (Vision)",
    "Audio Classification (Speech)",
    "ASR (Speech-to-Text)",
    "Text2Text Generation"
]
TASK_TO_PIPELINE_MAP = {
    "Causal Language Modeling (SFT/LoRA)": "text-generation",
    "DPO (Direct Preference Optimization)": "text-generation",
    "Question Answering (Text)": "question-answering",
    "Token Classification (NER)": "token-classification",
    "Sequence Classification (Text)": "text-classification",
    "Image Classification (Vision)": "image-classification",
    "Audio Classification (Speech)": "audio-classification",
    "ASR (Speech-to-Text)": "automatic-speech-recognition",
    "Text2Text Generation": "text2text-generation",
    "Text-to-Image Generation": "text-to-image",
}

MODEL_CARD_TEMPLATE = """---
language: es
license: apache-2.0
tags:
- autotrain-advanced
- fine-tuned
- {base_model_name}
widget:
- text: "Hola, ¿cómo estás?"
---
# {repo_id}
Este modelo es una versión afinada de [{base_model}](https://huggingface.co/{base_model}) entrenado con la herramienta [AutoTrain-Advanced](https://huggingface.co/spaces/autotrain-projects/autotrain-advanced).
## Detalles del Entrenamiento
- **Modo de Entrenamiento:** {training_mode}
- **Modelo Base:** `{base_model}`
- **Datasets:** `{datasets}`
- **Entrenado en:** {date}
### Hiperparámetros de Entrenamiento
```json
{hyperparameters}```
### Frameworks Utilizados
- Transformers
- PEFT
- Accelerate
- TRL
- Gradio
"""
DATASET_CARD_TEMPLATE = """---
license: mit
---
# {repo_id}
Este dataset fue creado utilizando la herramienta [AutoTrain-Advanced](https://huggingface.co/spaces/autotrain-projects/autotrain-advanced).
## Detalles del Dataset
- **Tipo de Creación:** {creation_type}
- **Modelo de Generación (si aplica):** `{generation_model}`
- **Fecha de Creación:** {date}
"""
_tox_pipe_singleton = None
@spaces.GPU
class DebiasingSFTTrainer(SFTTrainer):
    def __init__(self, *args, reweighting_terms=None, reweighting_factor=1.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.reweighting_terms = [term.strip().lower() for term in reweighting_terms] if reweighting_terms else []
        self.reweighting_factor = reweighting_factor

    def compute_loss(self, model, inputs, return_outputs=False):
        loss, outputs = super().compute_loss(model, inputs, return_outputs=True)
        if self.reweighting_terms and self.reweighting_factor > 1.0:
            input_ids = inputs.get("input_ids")
            decoded_texts = self.tokenizer.batch_decode(input_ids, skip_special_tokens=True)
            for text in decoded_texts:
                if any(term in text.lower() for term in self.reweighting_terms):
                    loss *= self.reweighting_factor
                    break
        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)
                    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

@spaces.GPU
def hf_login(token):
    if not token:
        return "Por favor, introduce un token."
    try:
        login(token=token, add_to_git_credential=True)
        user = whoami()
        return f"✅ Conectado como: {user['name']}"
    except Exception as e:
        return f"❌ Error en la conexión: {e}"
@spaces.GPU
def _clean_text(example, text_col, **kwargs):
    text = example.get(text_col, "")
    if not isinstance(text, str):
        return example
    if kwargs.get('remove_html_tags'):
        text = BeautifulSoup(text, "html.parser").get_text()
    if kwargs.get('remove_urls_emails'):
        text = re.sub(r'http\S+|www\S+|httpsS+', '', text, flags=re.MULTILINE)
    if kwargs.get('normalize_whitespace'):
        text = ' '.join(text.split())
    if kwargs.get('redact_pii'):
        text = re.sub(r'\S+@\S+', '<EMAIL>', text)
        text = re.sub(r'(\d{1,4}[-.\s]?){7,}|(\+\d{1,3}\s?)?\(?\d{3}\)?[\s.-]?\d{3}[\s.-]?\d{4}', '<PHONE>', text)
        text = re.sub(r'\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b', '<IP_ADDRESS>', text)
    example[text_col] = text
    return example
@spaces.GPU
def _apply_quality_filters(example, text_col, min_len, max_len, rep_threshold, exclude_keywords):
    text = example.get(text_col, "")
    if not isinstance(text, str): return False
    text_len = len(text.split())
    if not (min_len <= text_len <= max_len): return False
    words = text.split()
    if not words: return False
    word_counts = {}
    for word in words: word_counts[word] = word_counts.get(word, 0) + 1
    if not word_counts or (max(word_counts.values()) / len(words)) > rep_threshold: return False
    lower_text = text.lower()
    return not any(keyword in lower_text for keyword in exclude_keywords)
@spaces.GPU
def _apply_coherence_filter(example, text_col, char_rep_threshold, ngram_rep_threshold, entropy_threshold):
    """
    Filter out incoherent/gibberish text like:
    'οςords..header cover Korordsordsordsordsordsordsordsordsardordsardordsardoseotal...'
    
    Detects:
    - Excessive character repetition
    - Excessive n-gram repetition (substring patterns)
    - Low entropy (randomness)
    - Mixed/unusual character sequences
    """
    text = example.get(text_col, "")
    if not isinstance(text, str) or not text: 
        return False
    
    # Check for excessive character repetition (e.g., "aaaaaaa", "......")
    char_repetition_ratio = 0
    if len(text) > 0:
        for char in set(text):
            if char.isalnum() or char in '.,;:!?':
                char_count = text.count(char)
                char_ratio = char_count / len(text)
                char_repetition_ratio = max(char_repetition_ratio, char_ratio)
        
        if char_repetition_ratio > char_rep_threshold:
            return False
    
    # Check for excessive n-gram repetition (detect patterns like "ordsords", "WeekWeek")
    # Count how many distinct n-grams are repeated excessively
    text_lower = text.lower()
    
    # Count total characters that are part of highly repeated n-grams
    repeated_chars = 0
    ngram_counts = {}
    
    for n in [3, 4, 5]:
        if len(text_lower) >= n:
            for i in range(len(text_lower) - n + 1):
                ngram = text_lower[i:i+n]
                if ngram.isalpha():  # Only count alphabetic n-grams
                    ngram_counts[ngram] = ngram_counts.get(ngram, 0) + 1
    
    # Count characters covered by n-grams that appear more than 3 times
    if ngram_counts:
        highly_repeated_ngrams = {ng for ng, count in ngram_counts.items() if count > 3}
        if highly_repeated_ngrams:
            # Check what fraction of text is covered by these repeated patterns
            covered_positions = set()
            for i in range(len(text_lower)):
                for n in [3, 4, 5]:
                    if i + n <= len(text_lower):
                        ngram = text_lower[i:i+n]
                        if ngram in highly_repeated_ngrams:
                            for j in range(i, i+n):
                                covered_positions.add(j)
            
            repetition_coverage = len(covered_positions) / len(text_lower)
            if repetition_coverage > ngram_rep_threshold:
                return False
    
    # Check for low entropy (randomness) - coherent text has moderate entropy
    # Calculate character-level entropy
    if len(text) > 10:
        char_freq = {}
        for char in text:
            char_freq[char] = char_freq.get(char, 0) + 1
        
        entropy = 0
        for count in char_freq.values():
            p = count / len(text)
            if p > 0:
                entropy -= p * math.log2(p)
        
        # Normalize entropy by expected max entropy for this alphabet size
        max_entropy = math.log2(len(char_freq)) if len(char_freq) > 0 else 1
        normalized_entropy = entropy / max_entropy if max_entropy > 0 else 0
        
        # Too low entropy indicates repetitive/gibberish text
        if normalized_entropy < entropy_threshold:
            return False
    
    # Check for unusual character sequences (e.g., "}}_{", "/$")
    # Count ratio of alphanumeric vs special characters
    if len(text) > 0:
        alnum_count = sum(1 for c in text if c.isalnum() or c.isspace())
        alnum_ratio = alnum_count / len(text)
        # If less than 70% alphanumeric+space, likely gibberish
        if alnum_ratio < 0.7:
            return False
    
    # Check for mixed scripts (e.g., Greek + Latin)
    # This often indicates copy-paste errors or gibberish
    scripts = {
        'greek': sum(1 for c in text if '\u0370' <= c <= '\u03FF'),
        'cyrillic': sum(1 for c in text if '\u0400' <= c <= '\u04FF'),
        'arabic': sum(1 for c in text if '\u0600' <= c <= '\u06FF'),
        'chinese': sum(1 for c in text if '\u4E00' <= c <= '\u9FFF'),
    }
    non_latin_chars = sum(scripts.values())
    latin_chars = sum(1 for c in text if c.isalpha() and not any(scripts.values()))
    
    # If we have both significant Latin and non-Latin text, might be gibberish
    if non_latin_chars > 2 and latin_chars > 10:
        # Mixed scripts detected
        return False
    
    return True
@spaces.GPU
def _get_filter_functions(**kwargs):
    filters = []
    if kwargs.get('enable_quality_filter'):
        exclude_list = [k.strip().lower() for k in (kwargs.get('exclude_keywords_input', '') + ',' + kwargs.get('bias_keywords_input', '')).split(",") if k.strip()]
        filters.append(lambda ex: _apply_quality_filters(ex, kwargs['text_col'], kwargs['min_len_input'], kwargs['max_len_input'], kwargs['rep_threshold_input'], exclude_list))
    if kwargs.get('enable_language_filter'):
        allowed_langs = [lang.strip() for lang in kwargs.get('allowed_languages', 'en').split(',')]
        lang_threshold = kwargs.get('language_detection_threshold', 0.95)
        def lang_filter(ex):
            text = ex.get(kwargs['text_col'], "")
            if not text or not isinstance(text, str) or len(text.split()) < 5: return True
            try:
                detected = detect_langs(text)
                return any(lang.lang in allowed_langs and lang.prob > lang_threshold for lang in detected)
            except:
                return False
        filters.append(lang_filter)
    if kwargs.get('enable_toxicity_filter'):
        tox_threshold = kwargs.get('toxicity_threshold', 0.8)
        def tox_filter(ex):
            global _tox_pipe_singleton
            if _tox_pipe_singleton is None:
                logger.info("Initializing toxicity filter pipeline...")
                _tox_pipe_singleton = pipeline("text-classification", model="unitary/toxic-bert")
            text = ex.get(kwargs['text_col'], "")
            if not text or not isinstance(text, str): return True
            try:
                results = _tox_pipe_singleton(text[:512], truncation=True)
                return not (results[0]['label'] == 'toxic' and results[0]['score'] > tox_threshold)
            except Exception:
                return True
        filters.append(tox_filter)
    if kwargs.get('enable_coherence_filter'):
        char_rep_thresh = kwargs.get('coherence_char_repetition_threshold', 0.4)
        ngram_rep_thresh = kwargs.get('coherence_ngram_repetition_threshold', 0.3)
        entropy_thresh = kwargs.get('coherence_entropy_threshold', 0.5)
        filters.append(lambda ex: _apply_coherence_filter(ex, kwargs['text_col'], char_rep_thresh, ngram_rep_thresh, entropy_thresh))
    if any([kwargs.get('enable_readability_filter'), kwargs.get('enable_stopword_filter'), kwargs.get('enable_uniqueness_filter')]):
        # Basic English stop words list
        stop_words = set(['the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by', 'from', 'as', 'is', 'was', 'are', 'were', 'be', 'been', 'being', 'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'should', 'could', 'can', 'may', 'might', 'must', 'this', 'that', 'these', 'those', 'i', 'you', 'he', 'she', 'it', 'we', 'they', 'what', 'which', 'who', 'when', 'where', 'why', 'how'])
        def stats_filter(ex):
            text = ex.get(kwargs['text_col'], "")
            if not isinstance(text, str) or not text: return True
            words = text.split()
            num_words = len(words)
            if num_words == 0: return True
            if kwargs.get('enable_readability_filter'):
                try:
                    score = textstat.flesch_reading_ease(text)
                    if not (kwargs['min_readability'] <= score <= kwargs['max_readability']): return False
                except:
                    pass  # If readability calculation fails, keep the text
            if kwargs.get('enable_stopword_filter'):
                stopword_count = sum(1 for word in words if word.lower() in stop_words)
                if num_words > 0 and (stopword_count / num_words) > kwargs['max_stopword_ratio']: return False
            if kwargs.get('enable_uniqueness_filter'):
                if (len(set(words)) / num_words) < kwargs['min_uniqueness_ratio']: return False
            return True
        filters.append(stats_filter)
    return filters
@spaces.GPU
def _load_hf_streaming(ids, split="train", probabilities=None):
    streams = []
    valid_ids = []
    for ident in ids:
        try:
            d = load_dataset(ident, streaming=True, trust_remote_code=True)
            split_found = False
            if isinstance(d, dict):
                for s_name, ds in d.items():
                    if s_name.lower() == split or (split == "train" and "train" in s_name.lower()):
                        streams.append(ds)
                        split_found = True
                        break
            else:
                streams.append(d)
                split_found = True
            if split_found:
                valid_ids.append(ident)
            else:
                logger.warning(f"Split '{split}' not found in dataset {ident}. Excluding from this source.")
        except Exception as e:
            logger.error(f"Error loading dataset {ident} split {split}: {e}. Excluding from this source.")
    if not streams:
        return 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
    return interleave_datasets(streams, probabilities=probabilities)
@spaces.GPU
def _load_uploaded_stream(files):
    all_rows = []
    for f in files or []:
        content = f.read().decode("utf-8", errors="ignore")
        name = f.name.lower()
        if name.endswith(".csv"):
            import csv
            all_rows.extend(list(csv.DictReader(io.StringIO(content))))
        elif name.endswith(".jsonl"):
            all_rows.extend([json.loads(line) for line in io.StringIO(content) if line.strip()])
        elif name.endswith(".json"):
            data = json.loads(content)
            all_rows.extend(data if isinstance(data, list) else [data])
        elif name.endswith(".txt"):
            all_rows.extend([{"text": line} for line in io.StringIO(content) if line.strip()])
    if not all_rows:
        return None
    val_size = max(1, int(len(all_rows) * 0.01))
    random.shuffle(all_rows)
    return {"train": all_rows[:-val_size] if val_size > 0 else all_rows, "validation": all_rows[-val_size:] if val_size > 0 else []}
@spaces.GPU
def _guess_columns(sample):
    text_col, image_col, audio_col, label_col = "text", "image", "audio", "label"
    if not isinstance(sample, dict):
        return text_col, image_col, audio_col, label_col
    keys = {k.lower(): k for k in sample.keys()}
    if "text" in keys: text_col = keys["text"]
    elif "content" in keys: text_col = keys["content"]
    elif "prompt" in keys: text_col = keys["prompt"]
    if "image" in keys: image_col = keys["image"]
    elif "img" in keys: image_col = keys["img"]
    if "audio" in keys: audio_col = keys["audio"]
    elif "speech" in keys: audio_col = keys["speech"]
    if "label" in keys: label_col = keys["label"]
    elif "labels" in keys: label_col = keys["labels"]
    return text_col, image_col, audio_col, label_col
@spaces.GPU
def _apply_cda(dataset, text_col, cda_config_str):
    try:
        swap_groups = json.loads(cda_config_str)
    except (json.JSONDecodeError, ValueError) as e:
        logger.error(f"Configuración de CDA inválida: {e}.")
        return dataset
    def cda_generator():
        for example in dataset:
            original_text = example.get(text_col, "")
            if not isinstance(original_text, str):
                yield example
                continue
            yield example
            generated_texts = {original_text}
            current_texts = {original_text}
            for group in swap_groups:
                next_texts = set()
                for text in current_texts:
                    for word_to_replace in group:
                        if word_to_replace in text:
                            for replacement_word in group:
                                if word_to_replace != replacement_word:
                                    new_text = text.replace(word_to_replace, replacement_word)
                                    if new_text not in generated_texts:
                                        new_example = example.copy()
                                        new_example[text_col] = new_text
                                        yield new_example
                                        generated_texts.add(new_text)
                                        next_texts.add(new_text)
                current_texts.update(next_texts)
    return IterableDataset.from_generator(cda_generator)
@spaces.GPU
def _apply_back_translation(dataset, text_col, ratio, model_id, reverse_model_id):
    if not ratio or ratio <= 0:
        return dataset
    logger.info(f"Aplicando retrotraducción al {ratio*100}% del dataset.")
    try:
        pipe_to = pipeline("translation", model=model_id, )
        pipe_from = pipeline("translation", model=reverse_model_id, )
    except Exception as e:
        logger.error(f"No se pudieron cargar los modelos de traducción: {e}")
        return dataset
    def bt_generator():
        for example in dataset:
            yield example
            if random.random() < ratio:
                original_text = example.get(text_col, "")
                if isinstance(original_text, str) and original_text:
                    try:
                        translated = pipe_to(original_text, max_length=512)[0]['translation_text']
                        back_translated = pipe_from(translated, max_length=512)[0]['translation_text']
                        if back_translated:
                            new_example = example.copy()
                            new_example[text_col] = back_translated
                            yield new_example
                    except Exception as e:
                        logger.warning(f"Error en retrotraducción: {e}")
    return IterableDataset.from_generator(bt_generator)
@spaces.GPU
def _generate_synthetic_data(original_dataset, text_col, model_id, num_samples, prompt_template):
    if not num_samples or num_samples <= 0:
        return None
    logger.info(f"Iniciando generación de {num_samples} muestras sintéticas con el modelo {model_id}.")
    try:
        generator = pipeline("text-generation", model=model_id, )
    except Exception as e:
        logger.error(f"No se pudo cargar el modelo generador sintético: {e}")
        return None
    seed_examples = list(islice(original_dataset, 200))
    if not seed_examples:
        logger.warning("Dataset original vacío, no se pueden generar datos sintéticos.")
        return None
    def synthetic_generator():
        for i in range(num_samples):
            seed_example = random.choice(seed_examples)
            seed_text = seed_example.get(text_col, "")
            prompt = Template(prompt_template).render(example_text=seed_text)
            try:
                generated_output = generator(prompt, max_new_tokens=256, num_return_sequences=1, do_sample=True, temperature=0.9, top_p=0.95)
                cleaned_text = generated_output[0]['generated_text'][len(prompt):].strip()
                if "new example:" in cleaned_text.lower():
                   cleaned_text = re.split("new example:", cleaned_text, flags=re.IGNORECASE)[-1].strip()
                if cleaned_text:
                    new_example = seed_example.copy()
                    new_example[text_col] = cleaned_text
                    yield new_example
            except Exception as e:
                logger.warning(f"Error generando una muestra sintética: {e}")
                continue
    return IterableDataset.from_generator(synthetic_generator)
def _calculate_auto_config(block_size, is_gpt2_like, steps_per_epoch_estimate, batch_size, gradient_accumulation):
    safe_steps = int(steps_per_epoch_estimate or 10000)
    safe_batch_size = int(batch_size or 1)
    safe_grad_accum = int(gradient_accumulation or 8)
    safe_block_size = int(block_size or 1024)
    size = safe_steps * safe_batch_size * safe_grad_accum
    if size <= 1:
        size = 10000
    log_size = math.log2(max(1000, size))
    vocab_size = min(65536, 32000 + int(log_size * 2000))
    preliminary_hidden_size = max(512, min(4096, 512 + int(log_size * 100)))
    heads = max(8, min(32, preliminary_hidden_size // 64))
    if heads == 0: heads = 8
    hidden_size = (preliminary_hidden_size // heads) * heads
    layers = max(8, min(32, 8 + int(log_size * 1.5)))
    kv_heads = heads if is_gpt2_like else (max(1, heads // 4))
    return vocab_size, hidden_size, hidden_size * 2, layers, heads, safe_block_size, False, kv_heads
@spaces.GPU
def _get_eval_dataset(train_ds_id, eval_ds_id, uploaded_val_data, update_logs_fn):
    if eval_ds_id:
        yield update_logs_fn(f"Cargando dataset de evaluación: {eval_ds_id}", "Evaluación")
        return _load_hf_streaming([eval_ds_id], split="train")
    if uploaded_val_data:
        yield update_logs_fn("Usando split de validación de archivos subidos.", "Evaluación")
        return HFDataset.from_list(uploaded_val_data)
    if train_ds_id:
        yield update_logs_fn("Intentando cargar split 'validation' o 'test' del dataset de entrenamiento.", "Evaluación")
        try:
            for split_name in ["validation", "test"]:
                eval_ds = _load_hf_streaming([train_ds_id], split=split_name)
                if eval_ds:
                    yield update_logs_fn(f"Split '{split_name}' encontrado y cargado.", "Evaluación")
                    return eval_ds
        except Exception as e:
            yield update_logs_fn(f"Error cargando split de evaluación: {e}. Omitiendo.", "Evaluación")
            return None
    yield update_logs_fn("No se proporcionó dataset de evaluación. Omitiendo.", "Evaluación")
    return None
def _create_training_args(output_dir, repo_id, **kwargs):
    neftune_alpha = float(kwargs.get('neftune_noise_alpha', 0.0))
    optim_args_dict = {}
    if kwargs.get('optim_args'):
        try:
            optim_args_dict = ast.literal_eval(f"dict({kwargs['optim_args']})")
        except Exception as e:
            logger.warning(f"No se pudieron parsear los argumentos del optimizador: {e}.")
    args_dict = {
        "output_dir": os.path.join(output_dir, "results"),
        "per_device_train_batch_size": int(kwargs.get('batch_size', 1)),
        "gradient_accumulation_steps": int(kwargs.get('gradient_accumulation', 8)),
        "optim": kwargs.get('optimizer', 'adamw_torch'),
        "optim_args": optim_args_dict,
        "save_strategy": "steps",
        "logging_steps": int(kwargs.get('logging_steps', 10)),
        "save_steps": int(kwargs.get('save_steps', 50)),
        "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)),
        "max_grad_norm": float(kwargs.get('max_grad_norm', 1.0)),
        "warmup_ratio": float(kwargs.get('warmup_ratio', 0.03)),
        "lr_scheduler_type": kwargs.get('scheduler', 'cosine'),
        "weight_decay": float(kwargs.get('weight_decay', 0.01)),
        "load_best_model_at_end": kwargs.get('run_evaluation', False),
        "save_total_limit": int(kwargs.get('save_total_limit', 1)),
        "push_to_hub": True,
        "hub_model_id": repo_id,
        "hub_strategy": kwargs.get('hub_strategy', 'every_save'),
        "dataloader_num_workers": 2,
        "report_to": "wandb" if kwargs.get('wandb_api_key_input') else "none",
        "remove_unused_columns": False,
        "group_by_length": kwargs.get('group_by_length', False),
        "metric_for_best_model": kwargs.get('metric_for_best_model', 'loss') if kwargs.get('run_evaluation') else None,
        "greater_is_better": kwargs.get('greater_is_better', False),
        "neftune_noise_alpha": neftune_alpha if neftune_alpha > 0 else None,
        "adam_beta1": float(kwargs.get('adam_beta1', 0.9)),
        "adam_beta2": float(kwargs.get('adam_beta2', 0.999)),
        "adam_epsilon": float(kwargs.get('adam_epsilon', 1e-8)),
    }
    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

    max_steps_val = int(kwargs.get('max_steps', -1))
    if max_steps_val > 0:
        args_dict["max_steps"] = max_steps_val
    else:
        raise ValueError("Para datasets en streaming se requiere un valor positivo para 'Máximos Pasos de Entrenamiento'.")
    return TrainingArguments(**args_dict)
@spaces.GPU
def _generic_model_loader(model_name_or_path, model_class, **kwargs):
    config_kwargs = {"trust_remote_code": True}
    if kwargs.get('label2id'):
        config_kwargs.update({"label2id": kwargs['label2id'], "id2label": kwargs['id2label']})
    config = AutoConfig.from_pretrained(model_name_or_path, **config_kwargs)
    if kwargs.get('attention_dropout', 0) > 0: config.attention_dropout = kwargs['attention_dropout']
    if kwargs.get('hidden_dropout', 0) > 0: config.hidden_dropout = kwargs['hidden_dropout']
    model_kwargs = {
        "trust_remote_code": True,
        "config": config,
        "torch_dtype": torch.float32,
    }
    if kwargs.get('num_labels'):
        model_kwargs.update({"num_labels": kwargs['num_labels'], "ignore_mismatched_sizes": True})
    model = model_class.from_pretrained(model_name_or_path, **model_kwargs)
    return model
@spaces.GPU
def _find_all_linear_names(model):
    cls = torch.nn.Linear
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[-1])
    if 'lm_head' in lora_module_names:
        lora_module_names.remove('lm_head')
    common_targets = {'q_proj', 'v_proj', 'k_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'}
    return list(lora_module_names.intersection(common_targets)) or list(lora_module_names)
@spaces.GPU
def _sft_formatting_func(example, text_col, tokenizer, **kwargs):
    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.get(kwargs.get('tool_use_col_input', 'tools'))}</tool_code>")
        if example.get(kwargs.get('response_col_input', 'response')):
            response_parts.append(example.get(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 aplicando la plantilla de chat: {e}.")
                return "\n".join([m['content'] for m in messages])
        return ""
    return example.get(text_col, "")
@spaces.GPU
def _dpo_formatting_func(example, **kwargs):
    return {"prompt": example.get(kwargs.get('prompt_col_input', 'prompt'), ""), "chosen": example.get(kwargs.get('dpo_chosen_col_input', 'chosen'), ""), "rejected": example.get(kwargs.get('dpo_rejected_col_input', 'rejected'), "")}
@spaces.GPU
def _evaluate_perplexity(model, tokenizer, eval_dataset, text_col):
    model.eval()
    encodings = tokenizer("\n\n".join(ex[text_col] for ex in islice(eval_dataset, 1000)), return_tensors="pt")
    max_length = model.config.max_position_embeddings
    stride = 512
    seq_len = encodings.input_ids.size(1)
    nlls = []
    prev_end_loc = 0
    with torch.no_grad():
        for begin_loc in range(0, seq_len, stride):
            end_loc = min(begin_loc + max_length, seq_len)
            trg_len = end_loc - prev_end_loc
            input_ids = encodings.input_ids[:, begin_loc:end_loc]
            target_ids = input_ids.clone()
            target_ids[:, :-trg_len] = -100
            outputs = model(input_ids, labels=target_ids)
            neg_log_likelihood = outputs.loss
            nlls.append(neg_log_likelihood)
            prev_end_loc = end_loc
            if end_loc == seq_len:
                break
    ppl = torch.exp(torch.stack(nlls).mean())
    return ppl.item()
@spaces.GPU
def _merge_multiple_loras(base_model_id, adapter_ids_str, weights_str, combination_type):
    adapter_ids = [s.strip() for s in adapter_ids_str.split(',') if s.strip()]
    if not adapter_ids:
        yield "No se proporcionaron IDs de adaptadores válidos. Omitiendo la fusión múltiple."
        return base_model_id
    try:
        weights = [float(w.strip()) for w in weights_str.split(',')]
    except:
        weights = [1.0] * len(adapter_ids)
    if len(weights) != len(adapter_ids):
        weights = [1.0] * len(adapter_ids)
        yield "Pesos de adaptadores inválidos, usando 1.0 para todos."
    yield f"Cargando modelo base {base_model_id} para fusión múltiple..."
    model = AutoModelForCausalLM.from_pretrained(base_model_id, torch_dtype=torch.float32, trust_remote_code=True, )
    for i, adapter_id in enumerate(adapter_ids):
        yield f"Cargando adaptador {i+1}: {adapter_id}"
        model.load_adapter(adapter_id, adapter_name=f"adapter_{i}")
    adapter_names = [f"adapter_{i}" for i in range(len(adapter_ids))]
    yield f"Combinando adaptadores: {adapter_names} con pesos: {weights} y tipo: {combination_type}"
    model.add_weighted_adapter(adapters=adapter_names, weights=weights, adapter_name="combined", combination_type=combination_type)
    model.set_adapter("combined")
    yield "Fusionando combinación de adaptadores en el modelo base..."
    merged_model = model.merge_and_unload()
    temp_dir = tempfile.mkdtemp()
    yield f"Guardando modelo fusionado en {temp_dir}"
    merged_model.save_pretrained(temp_dir)
    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
    tokenizer.save_pretrained(temp_dir)
    yield f"Fusión de adaptadores completada. El entrenamiento continuará con el modelo fusionado en {temp_dir}."
    return temp_dir
@spaces.GPU
def _run_trainer_and_upload(trainer, tokenizer, repo_id, update_logs_fn, model_card_content, **kwargs):
    yield update_logs_fn("Iniciando ciclo de entrenamiento...", "Entrenando")
    trainer.train(resume_from_checkpoint=kwargs.get('resume_from_checkpoint') or False)
    final_metrics = {}
    if kwargs.get('run_evaluation'):
        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
    trainer.save_model(output_dir)
    if tokenizer:
        tokenizer.save_pretrained(output_dir)
    with open(os.path.join(output_dir, "README.md"), "w", encoding="utf-8") as f:
        f.write(model_card_content)
    yield update_logs_fn("Subiendo al Hub...", "Subiendo")
    upload_folder(folder_path=output_dir, repo_id=repo_id, commit_message="Fin de entrenamiento")
    # Don't explicitly delete trainer or call gc.collect() to avoid AcceleratorState issues
    # Let Python's garbage collector handle cleanup naturally
    return output_dir, final_metrics
@spaces.GPU
def train_sft_dpo(model_name, train_dataset, repo_id, update_logs_fn, model_card_content, **kwargs):
    output_dir = tempfile.mkdtemp()
    is_dpo = kwargs.get('training_mode') == "DPO (Direct Preference Optimization)"
    text_col = kwargs.get('text_col')
    try:
        tokenizer_id = kwargs.get('tokenizer_name_input') or model_name
        yield update_logs_fn(f"Cargando tokenizer '{tokenizer_id}'...", "Configuración")
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, 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']
        yield update_logs_fn(f"Cargando modelo '{model_name}'...", "Configuración")
        model = _generic_model_loader(model_name, AutoModelForCausalLM, **kwargs)
        peft_config = None
        if kwargs.get('peft'):
            target_modules = kwargs.get('target_modules').split(",") if not kwargs.get('auto_find_target_modules') else _find_all_linear_names(model)
            yield update_logs_fn(f"Módulos LoRA detectados/especificados: {target_modules}", "Configuración")
            peft_config = LoraConfig(
                r=int(kwargs.get('lora_r')), lora_alpha=int(kwargs.get('lora_alpha')), lora_dropout=float(kwargs.get('lora_dropout')),
                target_modules=target_modules, bias="none", task_type="CAUSAL_LM", use_dora=kwargs.get('use_dora', False),
                use_rslora=kwargs.get('use_rslora', False), init_lora_weights=kwargs.get('init_lora_weights', 'gaussian'),
                modules_to_save=kwargs.get('modules_to_save').split(',') if kwargs.get('modules_to_save') else None
            )
        training_args = _create_training_args(output_dir, repo_id, **kwargs)
        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):
                    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}
        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()
            if 'text_col' in sft_kwargs:
                del sft_kwargs['text_col']
            trainer_kwargs.update({"formatting_func": lambda ex: _sft_formatting_func(ex, text_col=text_col, tokenizer=tokenizer, **sft_kwargs)})
            if kwargs.get('enable_loss_reweighting'):
                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)
        return final_model_path, final_metrics
    except Exception as e:
        raise Exception(f"Error en {'DPO' if is_dpo else 'SFT'}: {e}\n{traceback.format_exc()}")
@spaces.GPU
def train_sequence_classification(model_name, train_dataset, repo_id, update_logs_fn, model_card_content, **kwargs):
    output_dir = tempfile.mkdtemp()
    try:
        labels = [s.strip() for s in kwargs['classification_labels'].split(',')]
        label2id = {l: i for i, l in enumerate(labels)}
        id2label = {i: l for i, l in enumerate(labels)}
        tokenizer_id = kwargs.get('tokenizer_name_input') 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)
        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):
                    yield update
                else:
                    eval_dataset = update
            if eval_dataset: eval_dataset = eval_dataset.map(preprocess, batched=True)
        metric = hf_evaluate.load("accuracy")
        def compute_metrics(eval_pred):
            logits, labels = eval_pred
            predictions = np.argmax(logits, axis=-1)
            return metric.compute(predictions=predictions, references=labels)
        training_args = _create_training_args(output_dir, repo_id, **kwargs)
        trainer = Trainer(
            model=model, args=training_args, train_dataset=train_dataset,
            eval_dataset=eval_dataset, compute_metrics=compute_metrics,
            tokenizer=tokenizer, data_collator=DataCollatorWithPadding(tokenizer=tokenizer)
        )
        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
    except Exception as e:
        raise Exception(f"Error en Sequence Classification: {e}\n{traceback.format_exc()}")
@spaces.GPU
def train_token_classification(model_name, train_dataset, repo_id, update_logs_fn, model_card_content, **kwargs):
    output_dir = tempfile.mkdtemp()
    try:
        labels = [s.strip() for s in kwargs['classification_labels'].split(',')]
        label2id = {l: i for i, l in enumerate(labels)}
        id2label = {i: l for i, l in enumerate(labels)}
        tokenizer_id = kwargs.get('tokenizer_name_input') or model_name
        yield update_logs_fn(f"Cargando tokenizer '{tokenizer_id}'...", "Configuración")
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True, add_prefix_space=True)
        yield update_logs_fn(f"Cargando modelo '{model_name}'...", "Configuración")
        model = _generic_model_loader(model_name, AutoModelForTokenClassification, num_labels=len(labels), label2id=label2id, id2label=id2label, **kwargs)
        def tokenize_and_align_labels(examples):
            tokenized_inputs = tokenizer(examples["tokens"], truncation=True, is_split_into_words=True)
            labels = []
            for i, label in enumerate(examples["ner_tags"]):
                word_ids = tokenized_inputs.word_ids(batch_index=i)
                previous_word_idx = None
                label_ids = []
                for word_idx in word_ids:
                    if word_idx is None or word_idx == previous_word_idx:
                        label_ids.append(-100)
                    else:
                        label_ids.append(label[word_idx])
                    previous_word_idx = word_idx
                labels.append(label_ids)
            tokenized_inputs["labels"] = labels
            return tokenized_inputs
        train_dataset = train_dataset.map(tokenize_and_align_labels, batched=True)
        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):
                    yield update
                else:
                    eval_dataset = update
            if eval_dataset: eval_dataset = eval_dataset.map(tokenize_and_align_labels, batched=True)
        metric = hf_evaluate.load("seqeval")
        def compute_metrics(p):
            predictions, labels = p
            predictions = np.argmax(predictions, axis=2)
            true_predictions = [[id2label[p] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels)]
            true_labels = [[id2label[l] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels)]
            results = metric.compute(predictions=true_predictions, references=true_labels)
            return {"precision": results["overall_precision"], "recall": results["overall_recall"], "f1": results["overall_f1"], "accuracy": results["overall_accuracy"]}
        training_args = _create_training_args(output_dir, repo_id, **kwargs)
        trainer = Trainer(
            model=model, args=training_args, train_dataset=train_dataset,
            eval_dataset=eval_dataset, tokenizer=tokenizer,
            data_collator=DataCollatorForTokenClassification(tokenizer=tokenizer),
            compute_metrics=compute_metrics
        )
        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
    except Exception as e:
        raise Exception(f"Error en Token Classification: {e}\n{traceback.format_exc()}")
@spaces.GPU
def train_question_answering(model_name, train_dataset, repo_id, update_logs_fn, model_card_content, **kwargs):
    output_dir = tempfile.mkdtemp()
    try:
        tokenizer_id = kwargs.get('tokenizer_name_input') or model_name
        yield update_logs_fn(f"Cargando tokenizer '{tokenizer_id}'...", "Configuración")
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True)
        yield update_logs_fn(f"Cargando modelo '{model_name}'...", "Configuración")
        model = _generic_model_loader(model_name, AutoModelForQuestionAnswering, **kwargs)
        max_length = 384
        doc_stride = 128
        def prepare_train_features(examples):
            tokenized_examples = tokenizer(
                examples["question"],
                examples["context"],
                truncation="only_second",
                max_length=max_length,
                stride=doc_stride,
                return_overflowing_tokens=True,
                return_offsets_mapping=True,
                padding="max_length",
            )
            sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
            offset_mapping = tokenized_examples.pop("offset_mapping")
            tokenized_examples["start_positions"] = []
            tokenized_examples["end_positions"] = []
            for i, offsets in enumerate(offset_mapping):
                input_ids = tokenized_examples["input_ids"][i]
                cls_index = input_ids.index(tokenizer.cls_token_id)
                sequence_ids = tokenized_examples.sequence_ids(i)
                sample_index = sample_mapping[i]
                answers = examples["answers"][sample_index]
                if len(answers["answer_start"]) == 0:
                    tokenized_examples["start_positions"].append(cls_index)
                    tokenized_examples["end_positions"].append(cls_index)
                else:
                    start_char = answers["answer_start"][0]
                    end_char = start_char + len(answers["text"][0])
                    token_start_index = 0
                    while sequence_ids[token_start_index] != 1:
                        token_start_index += 1
                    token_end_index = len(input_ids) - 1
                    while sequence_ids[token_end_index] != 1:
                        token_end_index -= 1
                    if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):
                        tokenized_examples["start_positions"].append(cls_index)
                        tokenized_examples["end_positions"].append(cls_index)
                    else:
                        while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
                            token_start_index += 1
                        tokenized_examples["start_positions"].append(token_start_index - 1)
                        while offsets[token_end_index][1] >= end_char:
                            token_end_index -= 1
                        tokenized_examples["end_positions"].append(token_end_index + 1)
            return tokenized_examples
        train_dataset = train_dataset.map(prepare_train_features, batched=True, remove_columns=next(iter(train_dataset)).keys())
        eval_dataset = None
        if kwargs.get('run_evaluation'):
            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):
                    yield update
                else:
                    eval_dataset_raw = update
            if eval_dataset_raw:
                eval_dataset = eval_dataset_raw.map(prepare_train_features, batched=True, remove_columns=next(iter(eval_dataset_raw)).keys())
        training_args = _create_training_args(output_dir, repo_id, **kwargs)
        data_collator = DefaultDataCollator()
        trainer = Trainer(
            model=model, args=training_args, train_dataset=train_dataset,
            eval_dataset=eval_dataset, tokenizer=tokenizer, data_collator=data_collator
        )
        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
    except Exception as e:
        raise Exception(f"Error en Question Answering: {e}\n{traceback.format_exc()}")
@spaces.GPU
def train_seq2seq(model_name, train_dataset, repo_id, update_logs_fn, model_card_content, **kwargs):
    output_dir = tempfile.mkdtemp()
    try:
        tokenizer_id = kwargs.get('tokenizer_name_input') or model_name
        yield update_logs_fn(f"Cargando tokenizer '{tokenizer_id}'...", "Configuración")
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, trust_remote_code=True)
        yield update_logs_fn(f"Cargando modelo '{model_name}'...", "Configuración")
        model = _generic_model_loader(model_name, AutoModelForSeq2SeqLM, **kwargs)
        def preprocess_function(examples):
            inputs = [ex[kwargs['text_col']] for ex in examples["translation"]]
            targets = [ex[kwargs['label_col']] for ex in examples["translation"]]
            model_inputs = tokenizer(inputs, max_length=128, truncation=True)
            with tokenizer.as_target_tokenizer():
                labels = tokenizer(targets, max_length=128, truncation=True)
            model_inputs["labels"] = labels["input_ids"]
            return model_inputs
        train_dataset = train_dataset.map(preprocess_function, batched=True)
        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):
                    yield update
                else:
                    eval_dataset = update
            if eval_dataset: eval_dataset = eval_dataset.map(preprocess_function, batched=True)
        metric = hf_evaluate.load("sacrebleu")
        def compute_metrics(eval_preds):
            preds, labels = eval_preds
            if isinstance(preds, tuple): preds = preds[0]
            decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
            decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
            decoded_preds = [pred.strip() for pred in decoded_preds]
            decoded_labels = [[label.strip()] for label in decoded_labels]
            result = metric.compute(predictions=decoded_preds, references=decoded_labels)
            return {"bleu": result["score"]}
        training_args_dict = _create_training_args(output_dir, repo_id, **kwargs).to_dict()
        training_args_dict["predict_with_generate"] = True
        training_args = Seq2SeqTrainingArguments(**training_args_dict)
        trainer = Seq2SeqTrainer(
            model=model, args=training_args, train_dataset=train_dataset,
            eval_dataset=eval_dataset, tokenizer=tokenizer,
            data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model),
            compute_metrics=compute_metrics
        )
        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
    except Exception as e:
        raise Exception(f"Error en Seq2Seq: {e}\n{traceback.format_exc()}")
@spaces.GPU
def train_text_to_image(model_name, train_dataset, repo_id, update_logs, model_card_content, **kwargs):
    """
    Entrena un modelo de difusión Text-to-Image usando stable diffusion.
    El dataset debe contener columnas 'text' (prompts) e 'image'.
    """
    output_dir = tempfile.mkdtemp()
    try:
        # Configuración inicial
        yield update_logs(f"Iniciando entrenamiento Text-to-Image con modelo base '{model_name}'...", "Configuración")
        
        # Cargar componentes del modelo
        from transformers import CLIPTextModel, CLIPTokenizer
        
        yield update_logs("Cargando componentes del modelo de difusión...", "Configuración")
        
        # Cargar tokenizer y text encoder
        tokenizer = CLIPTokenizer.from_pretrained(model_name, subfolder="tokenizer")
        text_encoder = CLIPTextModel.from_pretrained(model_name, subfolder="text_encoder")
        
        # Cargar VAE
        vae = AutoencoderKL.from_pretrained(model_name, subfolder="vae")
        
        # Cargar UNet (el componente que entrenaremos)
        unet = UNet2DConditionModel.from_pretrained(model_name, subfolder="unet")
        
        # Cargar scheduler
        noise_scheduler = DDPMScheduler.from_pretrained(model_name, subfolder="scheduler")
        
        yield update_logs("Componentes del modelo cargados exitosamente.", "Configuración")
        
        # Congelar VAE y text encoder
        vae.requires_grad_(False)
        text_encoder.requires_grad_(False)
        
        # Solo entrenaremos el UNet
        unet.train()
        
        # Configurar optimizador
        learning_rate = float(kwargs.get('learning_rate', 1e-5))
        optimizer = torch.optim.AdamW(
            unet.parameters(),
            lr=learning_rate,
            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))
        )
        
        yield update_logs("Optimizador configurado.", "Configuración")
        
        # Preparar dataset
        text_col = kwargs.get('text_col', 'text')
        image_col = kwargs.get('image_col', 'image')
        
        # Transformaciones para imágenes
        image_transforms = transforms.Compose([
            transforms.Resize(512, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(512),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ])
        
        def preprocess_train(examples):
            """Preprocesa ejemplos del dataset"""
            images = [image.convert("RGB") for image in examples[image_col]]
            examples["pixel_values"] = [image_transforms(image) for image in images]
            examples["input_ids"] = tokenizer(
                examples[text_col],
                padding="max_length",
                max_length=tokenizer.model_max_length,
                truncation=True,
                return_tensors="pt"
            ).input_ids
            return examples
        
        yield update_logs("Preprocesando dataset...", "Datos")
        train_dataset = train_dataset.map(preprocess_train, batched=True, remove_columns=[image_col])
        
        # Configuración de entrenamiento
        batch_size = int(kwargs.get('batch_size', 1))
        gradient_accumulation_steps = int(kwargs.get('gradient_accumulation', 4))
        max_steps = int(kwargs.get('max_steps', 1000))
        num_epochs = int(kwargs.get('num_epochs', 1))
        
        # Crear dataloader
        train_dataloader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=2
        )
        
        # Scheduler de learning rate
        from diffusers.optimization import get_scheduler as get_diffusers_lr_scheduler
        lr_scheduler = get_diffusers_lr_scheduler(
            kwargs.get('scheduler', 'cosine'),
            optimizer=optimizer,
            num_warmup_steps=int(max_steps * float(kwargs.get('warmup_ratio', 0.03))),
            num_training_steps=max_steps
        )
        
        yield update_logs(f"Iniciando entrenamiento: {max_steps} pasos, batch_size={batch_size}", "Entrenando")
        
        # Mover modelos a GPU si está disponible
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        unet = unet.to(device)
        vae = vae.to(device)
        text_encoder = text_encoder.to(device)
        
        # Loop de entrenamiento
        global_step = 0
        progress_bar = tqdm(range(max_steps), desc="Entrenando")
        
        for epoch in range(num_epochs):
            for step, batch in enumerate(train_dataloader):
                if global_step >= max_steps:
                    break
                
                # Convertir imágenes a latents usando VAE
                pixel_values = torch.stack(batch["pixel_values"]).to(device)
                with torch.no_grad():
                    latents = vae.encode(pixel_values).latent_dist.sample()
                    latents = latents * vae.config.scaling_factor
                
                # Añadir ruido a los latents
                noise = torch.randn_like(latents)
                timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=device).long()
                noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
                
                # Obtener embeddings del texto
                input_ids = batch["input_ids"].to(device)
                with torch.no_grad():
                    encoder_hidden_states = text_encoder(input_ids)[0]
                
                # Predecir el ruido
                noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
                
                # Calcular pérdida
                loss = F.mse_loss(noise_pred, noise, reduction="mean")
                
                # Backpropagation
                loss = loss / gradient_accumulation_steps
                loss.backward()
                
                if (step + 1) % gradient_accumulation_steps == 0:
                    # Gradient clipping
                    torch.nn.utils.clip_grad_norm_(unet.parameters(), float(kwargs.get('max_grad_norm', 1.0)))
                    
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
                    
                    global_step += 1
                    progress_bar.update(1)
                    
                    # Logging
                    if global_step % int(kwargs.get('logging_steps', 10)) == 0:
                        yield update_logs(f"Paso {global_step}/{max_steps} - Loss: {loss.item():.4f}", "Entrenando")
                    
                    # Guardar checkpoint
                    if global_step % int(kwargs.get('save_steps', 500)) == 0:
                        yield update_logs(f"Guardando checkpoint en paso {global_step}...", "Guardando")
                        checkpoint_dir = os.path.join(output_dir, f"checkpoint-{global_step}")
                        os.makedirs(checkpoint_dir, exist_ok=True)
                        unet.save_pretrained(os.path.join(checkpoint_dir, "unet"))
                        
                        # Si hub_strategy es every_save, subir checkpoint
                        if kwargs.get('hub_strategy') == 'every_save':
                            try:
                                upload_folder(
                                    folder_path=checkpoint_dir,
                                    repo_id=repo_id,
                                    commit_message=f"Checkpoint paso {global_step}"
                                )
                            except Exception as e:
                                yield update_logs(f"Advertencia: No se pudo subir checkpoint: {e}", "Guardando")
                    
                    if global_step >= max_steps:
                        break
            
            if global_step >= max_steps:
                break
        
        progress_bar.close()
        yield update_logs("Entrenamiento completado. Guardando modelo final...", "Guardando")
        
        # Guardar modelo final
        final_output_dir = os.path.join(output_dir, "final_model")
        os.makedirs(final_output_dir, exist_ok=True)
        
        # Crear pipeline completo
        pipeline = StableDiffusionText2ImagePipeline(
            text_encoder=text_encoder,
            vae=vae,
            unet=unet,
            tokenizer=tokenizer,
            scheduler=noise_scheduler,
            safety_checker=None,
            feature_extractor=None,
            requires_safety_checker=False
        )
        
        # Guardar pipeline completo
        pipeline.save_pretrained(final_output_dir)
        
        # Guardar model card
        with open(os.path.join(final_output_dir, "README.md"), "w", encoding="utf-8") as f:
            f.write(model_card_content)
        
        yield update_logs("Modelo guardado. Subiendo al Hub...", "Subiendo")
        
        # Subir al Hub
        upload_folder(
            folder_path=final_output_dir,
            repo_id=repo_id,
            commit_message="Entrenamiento Text-to-Image completado"
        )
        
        yield update_logs(f"✅ Modelo subido exitosamente a {repo_id}", "Completado")
        
        # Métricas finales
        final_metrics = {
            "final_loss": loss.item(),
            "total_steps": global_step,
            "epochs_completed": epoch + 1
        }
        
        # Limpieza
        del unet, vae, text_encoder, pipeline
        gc.collect()
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        
        return final_output_dir, final_metrics
        
    except Exception as e:
        yield update_logs(f"❌ Error en entrenamiento Text-to-Image: {str(e)}", "Error")
        raise Exception(f"Error en Text-to-Image: {e}\n{traceback.format_exc()}")
@spaces.GPU
def _get_data_processing_pipeline(**kwargs):
    hf_ids = [x.strip() for x in (kwargs.get('datasets_hf_text') or "").split(",") if x.strip()]
    if not hf_ids and not kwargs.get('uploads'):
        raise ValueError("No se proporcionaron datasets.")
    dataset_weights_str = kwargs.get('dataset_weights', '')
    probabilities = None
    if dataset_weights_str:
        try:
            probabilities = [float(w.strip()) for w in dataset_weights_str.split(',')]
        except ValueError:
            probabilities = None
    train_dataset, uploaded_val_data = None, None
    if kwargs.get('uploads'):
        uploaded_data_map = _load_uploaded_stream(kwargs.get('uploads'))
        if uploaded_data_map and 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:
        hf_train_dataset = _load_hf_streaming(hf_ids, split="train", probabilities=probabilities if not train_dataset else None)
        if hf_train_dataset:
            if train_dataset is None:
                train_dataset = hf_train_dataset
            else:
                all_streams = [train_dataset, hf_train_dataset]
                all_probs = [0.5, 0.5]
                train_dataset = interleave_datasets(all_streams, probabilities=all_probs)
    if train_dataset is None:
        raise ValueError("No se pudieron cargar datos de entrenamiento válidos.")
    try:
        first_example = next(iter(train_dataset))
    except StopIteration:
        raise ValueError("El dataset de entrenamiento está vacío después del procesamiento.")
    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 ["Image Classification (Vision)", "Audio Classification (Speech)"]
    if is_text_task:
        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))
        filters = _get_filter_functions(**kwargs)
        if filters:
            for f in filters:
                train_dataset = train_dataset.filter(f)
        if kwargs.get('enable_back_translation'):
            train_dataset = _apply_back_translation(train_dataset, text_col, kwargs['bt_augmentation_ratio'], kwargs['bt_model_id'], kwargs['bt_reverse_model_id'])
        if kwargs.get('enable_synthetic_data'):
            synthetic_ds = _generate_synthetic_data(train_dataset, text_col, kwargs['synthetic_model_id'], int(kwargs['num_synthetic_samples']), kwargs['synthetic_prompt_template'])
            if synthetic_ds:
                train_dataset = interleave_datasets([train_dataset, synthetic_ds])
        if kwargs.get('enable_cda') and kwargs.get('cda_json_config'):
            train_dataset = _apply_cda(train_dataset, text_col, kwargs['cda_json_config'])
        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))
            )
    return train_dataset, kwargs
@spaces.GPU
def _train_and_upload(progress=gr.Progress(), **kwargs):
    logs, repo_link, final_model_path, final_metrics = "", "", None, {}
    progress(0, desc="Iniciando...")
    yield (
        "Iniciando...",
        "Inicio",
        "",
        gr.update(value=None),
        gr.update(value="Entrenando...", interactive=False),
        gr.update(visible=True)
    )
    def update_logs(new_msg, phase_msg):
        nonlocal logs, repo_link, final_metrics
        logs += f"[{phase_msg}] {new_msg}\n"
        progress(0, desc=f"[{phase_msg}] {new_msg}")
        return (
            logs,
            phase_msg,
            repo_link,
            gr.update(value=final_metrics if final_metrics else None)
        )
    try:
        yield update_logs("Verificando autenticación...", "Inicio") + (gr.update(), gr.update())
        user = whoami()
        username = user.get("name")
        if not username:
             raise ValueError("No se pudo obtener el nombre de usuario de Hugging Face. Por favor, verifica tu token.")
        model_name = kwargs.get('model_base_input', '').strip()
        if kwargs.get('enable_multi_adapter_merge'):
            temp_model_path = model_name
            lora_merge_generator = _merge_multiple_loras(model_name, kwargs['multi_adapter_model_ids'], kwargs['multi_adapter_weights'], kwargs['multi_adapter_combination_type'])
            try:
                while True:
                    status = next(lora_merge_generator)
                    yield update_logs(status, "Fusión Múltiple") + (gr.update(), gr.update())
            except StopIteration as e:
                temp_model_path = e.value
            model_name = temp_model_path
        repo_name_input = kwargs.get('repo_name_input', '').strip()
        if repo_name_input:
            repo_base = re.sub(r'[^a-zA-Z0-9_.-]+', '-', repo_name_input)
            repo_base = re.sub(r'^[.-]+|[.-]+$', '', repo_base)
        else:
            model_name_base = model_name.split('/')[-1] if model_name else "finetuned-model"
            sanitized_model_name_base = re.sub(r'[^a-zA-Z0-9_.-]+', '-', model_name_base)
            sanitized_model_name_base = re.sub(r'^[.-]+|[.-]+$', '', sanitized_model_name_base)
            repo_base = f"{sanitized_model_name_base}-{uuid.uuid4().hex[:6]}"
        if not repo_base:
            repo_base = f"autotrain-model-{uuid.uuid4().hex[:8]}"
        max_repo_base_len = 96 - (len(username) + 1)
        repo_base = repo_base[:max_repo_base_len]
        repo_id = f"{username}/{repo_base}"
        yield update_logs(f"Creando o verificando repositorio: '{repo_id}'", "Inicio") + (gr.update(), gr.update())
        create_repo(repo_id, exist_ok=True, private=kwargs.get('private_repo', False))
        repo_link = f"https://huggingface.co/{repo_id}"
        yield update_logs("Repositorio listo.", "Inicio") + (gr.update(), gr.update())
        base_model_id_for_training = model_name
        if kwargs.get('train_from_scratch'):
            yield update_logs("Preparando entrenamiento desde cero...", "Modelo Cero") + (gr.update(), gr.update())
            architecture = kwargs.get('scratch_architecture')
            if not architecture or architecture not in ARCHITECTURE_MAP:
                raise ValueError(f"Arquitectura '{architecture}' no es válida o no está soportada para entrenamiento desde cero. Opciones válidas: {list(ARCHITECTURE_MAP.keys())}")
            config_class, model_class = ARCHITECTURE_MAP[architecture]
            if kwargs.get('auto_config_scratch'):
                vocab_size, hidden_size, intermediate_size, layers, heads, block_size_val, tie_word_embeddings, kv_heads = _calculate_auto_config(kwargs.get('block_size'), architecture == "GPT2", kwargs.get('steps_per_epoch_estimate'), kwargs.get('batch_size'), kwargs.get('gradient_accumulation'))
                config = config_class(vocab_size=vocab_size, hidden_size=hidden_size, intermediate_size=intermediate_size, num_hidden_layers=layers, num_attention_heads=heads, num_key_value_heads=kv_heads, max_position_embeddings=block_size_val, tie_word_embeddings=tie_word_embeddings)
                model = model_class(config)
            elif kwargs.get('manual_config_scratch'):
                vocab_size = int(kwargs.get('scratch_vocab_size', 32000))
                hidden_size = int(kwargs.get('scratch_hidden_size', 1024))
                intermediate_size = int(kwargs.get('scratch_intermediate_size', 2048))
                layers = int(kwargs.get('scratch_layers', 8))
                heads = int(kwargs.get('scratch_heads', 8))
                kv_heads = int(kwargs.get('scratch_kv_heads', 8))
                block_size_val = int(kwargs.get('scratch_block_size', 1024))
                tie_word_embeddings = kwargs.get('scratch_tie_word_embeddings', False)
                config = config_class(vocab_size=vocab_size, hidden_size=hidden_size, intermediate_size=intermediate_size, num_hidden_layers=layers, num_attention_heads=heads, num_key_value_heads=kv_heads, max_position_embeddings=block_size_val, tie_word_embeddings=tie_word_embeddings)
                model = model_class(config)
            else:
                raise ValueError("Debe seleccionar auto-configuración o configuración manual para entrenar desde cero.")
            temp_model_dir = tempfile.mkdtemp()
            model.save_pretrained(temp_model_dir)
            tokenizer_id = kwargs.get('tokenizer_name_input') or SCRATCH_TOKENIZER_MAP.get(architecture, "gpt2")
            try:
                tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
                tokenizer.save_pretrained(temp_model_dir)
                yield update_logs(f"Tokenizer base '{tokenizer_id}' guardado para el modelo desde cero.", "Modelo Cero") + (gr.update(), gr.update())
            except Exception as e:
                 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['tokenizer_name'] = temp_model_dir
            yield update_logs(f"Modelo {architecture} inicializado en {temp_model_dir}.", "Modelo Cero") + (gr.update(), gr.update())
        yield update_logs("Procesando y cargando datasets...", "Datos") + (gr.update(), gr.update())
        train_dataset, kwargs = _get_data_processing_pipeline(**kwargs)
        yield update_logs(f"Columnas detectadas (texto: {kwargs['text_col']}, imagen: {kwargs['image_col']})", "Datos") + (gr.update(), gr.update())
        if kwargs.get('wandb_api_key_input'):
            os.environ["WANDB_API_KEY"] = kwargs['wandb_api_key_input']
            os.environ["WANDB_PROJECT"] = kwargs.get('wandb_project_input') or f"{repo_base}"
            os.environ["WANDB_LOG_MODEL"] = "checkpoint"
        model_card_content = MODEL_CARD_TEMPLATE.format(
            repo_id=repo_id, base_model=model_name, base_model_name=_sanitize_model_name_for_yaml(model_name),
            training_mode=kwargs.get('training_mode'),
            datasets=', '.join([x.strip() for x in (kwargs.get('datasets_hf_text') or "").split(",") if x.strip()]) or "Archivos locales",
            hyperparameters=json.dumps({k: v for k, v in kwargs.items() if isinstance(v, (str, int, float, bool)) and 'token' not in k and 'key' not in k and v is not None}, indent=2),
            date=datetime.now().strftime("%Y-%m-%d")
        )
        training_mode = kwargs.get('training_mode')
        training_function_map = {
            "Causal Language Modeling (SFT/LoRA)": train_sft_dpo,
            "DPO (Direct Preference Optimization)": train_sft_dpo,
            "Question Answering (Text)": train_question_answering,
            "Sequence Classification (Text)": train_sequence_classification,
            "Token Classification (NER)": train_token_classification,
            "Text2Text Generation": train_seq2seq,
            "Text-to-Image Generation": train_text_to_image,
        }
        train_func = training_function_map.get(training_mode)
        if train_func:
            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
                except StopIteration as e:
                    final_model_path, final_metrics = e.value
                    break
        else:
            raise ValueError(f"El modo de entrenamiento '{training_mode}' no está implementado.")
        if kwargs.get('run_perplexity_evaluation') and final_model_path and training_mode in ["Causal Language Modeling (SFT/LoRA)", "DPO (Direct Preference Optimization)"]:
            yield update_logs("Iniciando evaluación de perplejidad...", "Evaluación Final") + (gr.update(), gr.update())
            model = AutoModelForCausalLM.from_pretrained(final_model_path, torch_dtype=torch.float32, )
            tokenizer = AutoTokenizer.from_pretrained(final_model_path)
            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):
                    yield update + (gr.update(), gr.update())
                else:
                    eval_dataset_perp = update
            if eval_dataset_perp:
                ppl = _evaluate_perplexity(model, tokenizer, eval_dataset_perp, kwargs['text_col'])
                final_metrics['perplexity'] = ppl
                yield update_logs(f"Evaluación de Perplejidad completada. Perplejidad: {ppl:.4f}", "Evaluación Final") + (gr.update(), gr.update())
        final_logs, final_phase, final_repo_link, _ = update_logs(f"✅ Entrenamiento y subida completados: {repo_link}", "Listo")
        yield (
            final_logs,
            final_phase,
            f"### ✅ [Modelo Finalizado: Visita el Repositorio en el Hub]({final_repo_link})",
            gr.update(value=final_metrics),
            gr.update(value="Iniciar Entrenamiento", interactive=True),
            gr.update(visible=False)
        )
    except Exception as e:
        err_msg = f"❌ Error fatal: {type(e).__name__}: {e}\n{traceback.format_exc()}"
        error_logs, error_phase, _, _ = update_logs(err_msg, "Error")
        yield (
            error_logs,
            error_phase,
            "",
            gr.update(value=None),
            gr.update(value="Iniciar Entrenamiento", interactive=True),
            gr.update(visible=False)
        )
@spaces.GPU
def run_inference(task_mode, model_id, text_in, context_in, image_in, audio_in, temperature, top_p, max_new_tokens):
    if not model_id: return "Por favor, introduce un ID de modelo del Hub.", model_id, gr.update(), gr.update(), gr.update(), gr.update()
    task_name = TASK_TO_PIPELINE_MAP.get(task_mode)
    if not task_name: return f"La inferencia para el modo '{task_mode}' no está soportada.", model_id, gr.update(), gr.update(), gr.update(), gr.update()
    try:
        pipe = pipeline(task_name, model=model_id, torch_dtype=torch.float32, trust_remote_code=True, )
        result = None
        if task_name == "text-generation":
            if not text_in: return "Por favor, introduce un prompt de texto.", model_id, gr.update(), gr.update(), gr.update(), gr.update()
            result = pipe(text_in, max_new_tokens=int(max_new_tokens), do_sample=True, temperature=temperature, top_p=top_p)
        elif task_name == "question-answering":
            if not text_in or not context_in: return "Por favor, introduce una pregunta y un contexto.", model_id, gr.update(), gr.update(), gr.update(), gr.update()
            result = pipe(question=text_in, context=context_in)
        elif task_name in ["token-classification", "text2text-generation", "text-classification"]:
            if not text_in: return f"Por favor, introduce texto para {task_name}.", model_id, gr.update(), gr.update(), gr.update(), gr.update()
            result = pipe(text_in)
        elif task_name in ["image-classification", "audio-classification", "automatic-speech-recognition"]:
            input_data = image_in if "image" in task_name else audio_in
            if input_data is None: return f"Por favor, proporciona una entrada de { 'imagen' if 'image' in task_name else 'audio' }.", model_id, gr.update(), gr.update(), gr.update(), gr.update()
            result = pipe(input_data)
        return f"Resultado:\n\n{json.dumps(result, indent=2, ensure_ascii=False)}", model_id, gr.update(), gr.update(), gr.update(), gr.update()
    except Exception as e: return f"Error en Inferencia: {e}\n{traceback.format_exc()}", model_id, gr.update(), gr.update(), gr.update(), gr.update()
def update_inference_ui(task_mode):
    task_name = TASK_TO_PIPELINE_MAP.get(task_mode, "")
    is_text_gen = task_name == "text-generation"
    show_text = task_name in ["text-generation", "text2text-generation", "token-classification", "question-answering", "text-classification", "text-to-image"]
    show_context = task_name == "question-answering"
    show_image = task_name in ["image-classification"]
    show_audio = task_name in ["audio-classification", "automatic-speech-recognition"]
    text_label = "Pregunta" if task_name == "question-answering" else "Entrada de Texto / Prompt"
    return (
        gr.update(visible=show_text, label=text_label),
        gr.update(visible=show_context),
        gr.update(visible=show_image),
        gr.update(visible=show_audio),
        gr.update(visible=is_text_gen)
    )
@spaces.GPU
def create_and_upload_dataset(hf_token, repo_name, creation_type, synth_model, synth_prompt, synth_num_samples, file_uploads, progress=gr.Progress()):
    if not hf_token:
        return "Error: Se requiere un token de Hugging Face.", ""
    if not repo_name:
        return "Error: Se requiere un nombre de repositorio para el dataset.", ""
    try:
        login(token=hf_token)
        user = whoami()
        username = user.get("name")
        repo_base = f"{username}-{uuid.uuid4().hex[:6]}" if not repo_name else re.sub(r'[^a-zA-Z0-9_.-]+', '-', repo_name)[:90]
        repo_id = f"{username}/{repo_base}"
        create_repo(repo_id, repo_type="dataset", exist_ok=True)
        all_data = []
        if creation_type == "Sintético":
            if not synth_model or not synth_prompt or not synth_num_samples:
                return "Error: Para la generación sintética se requiere un modelo, un prompt y un número de muestras.", ""
            progress(0, desc="Cargando modelo generador...")
            generator = pipeline("text-generation", model=synth_model, )
            for i in progress.tqdm(range(int(synth_num_samples)), desc="Generando muestras"):
                try:
                    generated_output = generator(synth_prompt, max_new_tokens=256, num_return_sequences=1, do_sample=True, temperature=0.9, top_p=0.95)
                    cleaned_text = generated_output[0]['generated_text'][len(synth_prompt):].strip()
                    if cleaned_text:
                        all_data.append({"text": cleaned_text})
                except Exception as e:
                    logger.warning(f"Error al generar muestra {i}: {e}")
        elif creation_type == "Basado en Archivo":
            if not file_uploads:
                return "Error: Por favor, sube al menos un archivo.", ""
            progress(0.5, desc="Procesando archivos subidos...")
            file_data = _load_uploaded_stream(file_uploads)
            all_data = file_data.get("train", []) + file_data.get("validation", [])
        if not all_data:
            return "Error: No se generaron o procesaron datos.", ""
        progress(0.8, desc="Guardando y subiendo al Hub...")
        with tempfile.TemporaryDirectory() as temp_dir:
            data_file = os.path.join(temp_dir, "data.jsonl")
            with open(data_file, "w", encoding="utf-8") as f:
                for item in all_data:
                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
            readme_content = DATASET_CARD_TEMPLATE.format(
                repo_id=repo_id,
                creation_type=creation_type,
                generation_model=synth_model if creation_type == "Sintético" else "N/A",
                date=datetime.now().strftime("%Y-%m-%d")
            )
            readme_file = os.path.join(temp_dir, "README.md")
            with open(readme_file, "w", encoding="utf-8") as f:
                f.write(readme_content)
            api = HfApi()
            api.upload_folder(
                folder_path=temp_dir,
                repo_id=repo_id,
                repo_type="dataset",
                commit_message="Creación de dataset con AutoTrain-Advanced"
            )
        dataset_link = f"https://huggingface.co/datasets/{repo_id}"
        return f"✅ Dataset creado y subido exitosamente a {repo_id}", f"### ✅ [Dataset Disponible: Visita el Repositorio]({dataset_link})"
    except Exception as e:
        return f"❌ Error fatal durante la creación del dataset: {e}\n{traceback.format_exc()}", ""
@spaces.GPU
def gradio_train_wrapper(*args):
    kwargs = dict(zip(all_input_components_dict.keys(), args))
    yield from _train_and_upload(**kwargs)
@spaces.GPU
def gradio_preview_data_wrapper(*args):
    kwargs = dict(zip(all_input_components_dict.keys(), args))
    try:
        preview_text = "Procesando vista previa...\n"
        yield preview_text
        model_id_for_tokenizer = kwargs.get('model_base_input')
        if not model_id_for_tokenizer and not kwargs.get('train_from_scratch'):
            raise ValueError("Se necesita un ID de modelo base para cargar el tokenizer para la vista previa.")
        
        dataset, processed_kwargs = _get_data_processing_pipeline(**kwargs)
        text_col = processed_kwargs.get('text_col')

        if kwargs.get('train_from_scratch'):
            tokenizer_id = SCRATCH_TOKENIZER_MAP.get(kwargs.get('scratch_architecture'), 'gpt2')
        else:
            tokenizer_id = kwargs.get('tokenizer_name_input') or model_id_for_tokenizer
            
        tokenizer = AutoTokenizer.from_pretrained(
            tokenizer_id, 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']
        
        preview_samples = []
        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)
            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")
        
        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."
        yield preview_text
    except Exception as e:
        yield f"Error al generar la vista previa: {e}\n{traceback.format_exc()}"
def toggle_training_mode_ui(is_scratch):
    return (
        gr.update(visible=not is_scratch),  # model_base_input
        gr.update(visible=not is_scratch),  # tokenizer_name_input
        gr.update(visible=not is_scratch),  # multi_adapter_accordion
        gr.update(visible=not is_scratch),  # peft_accordion
        gr.update(visible=is_scratch),      # auto_config_scratch
        gr.update(visible=is_scratch),      # scratch_architecture
        gr.update(visible=is_scratch),      # manual_config_scratch
        gr.update(visible=is_scratch),      # scratch_vocab_size
        gr.update(visible=is_scratch),      # scratch_hidden_size
        gr.update(visible=is_scratch),      # scratch_intermediate_size
        gr.update(visible=is_scratch),      # scratch_layers
        gr.update(visible=is_scratch),      # scratch_heads
        gr.update(visible=is_scratch),      # scratch_kv_heads
        gr.update(visible=is_scratch),      # scratch_block_size
        gr.update(visible=is_scratch),      # scratch_tie_word_embeddings
        gr.update(visible=is_scratch),      # steps_per_epoch_estimate
        gr.update(visible=is_scratch),      # attention_dropout
        gr.update(visible=is_scratch),      # hidden_dropout
    )
def toggle_task_specific_ui(training_mode):
    is_classification = "Classification" in training_mode
    is_dpo = "DPO" in training_mode
    is_sft = "Causal" in training_mode
    is_ner = "Token Classification" in training_mode
    is_diffusion = "Image Generation" in training_mode
    return (
        gr.update(visible=is_classification or is_ner),
        gr.update(visible=is_dpo),
        gr.update(visible=is_sft),
        gr.update(visible=is_diffusion),
        gr.update(visible=not is_diffusion)
    )
def toggle_sft_format_ui(format_style):
    is_tool = format_style == "Razonamiento/Herramientas"
    return gr.update(visible=is_tool)
def toggle_auto_modules_ui(is_auto):
    return gr.update(visible=not is_auto)
def toggle_dataset_creator_ui(choice):
    is_synth = choice == "Sintético"
    return gr.update(visible=is_synth), gr.update(visible=not is_synth)

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
    gr.Markdown("# 🚀 AutoTrain-Advanced: Tu Plataforma de Entrenamiento de Modelos")
    gr.Markdown("### Una interfaz completa para fine-tuning y PEFT (LoRA).")

    with gr.Tab("1. Autenticación"):
        gr.Markdown("#### Conecta tu cuenta de Hugging Face para guardar y cargar modelos.")
        with gr.Row():
            hf_token_input = gr.Textbox(label="Token de Hugging Face (con permisos de escritura)", type="password", placeholder="hf_...", scale=3)
            login_button = gr.Button("Conectar", variant="primary", scale=1)
        login_status = gr.Textbox(label="Estado de Conexión", interactive=False)
        login_button.click(hf_login, inputs=[hf_token_input], outputs=[login_status])

    with gr.Tab("2. Creación de Dataset"):
        gr.Markdown("## 🧩 Genera o Procesa Datasets y Súbelos al Hub")
        with gr.Row():
            with gr.Column(scale=1):
                dset_repo_name = gr.Textbox(label="Nombre del Repositorio del Dataset", placeholder="mi-nuevo-dataset")
                dset_creation_type = gr.Radio(["Sintético", "Basado en Archivo"], label="Tipo de Creación", value="Sintético")
                with gr.Group(visible=True) as dset_synth_group:
                    dset_synth_model = gr.Textbox(label="Modelo Generador", placeholder="p.ej. 'mistralai/Mistral-7B-Instruct-v0.2'")
                    dset_synth_prompt = gr.Textbox(label="Prompt de Generación", lines=5, placeholder="Escribe una reseña de producto de 5 estrellas para...")
                    dset_synth_num_samples = gr.Number(label="Número de Muestras", value=100)
                with gr.Group(visible=False) as dset_file_group:
                    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_link_output = gr.Markdown()
        dset_creation_type.change(toggle_dataset_creator_ui, inputs=[dset_creation_type], outputs=[dset_synth_group, dset_file_group])
        dset_create_button.click(
            create_and_upload_dataset,
            inputs=[hf_token_input, dset_repo_name, dset_creation_type, dset_synth_model, dset_synth_prompt, dset_synth_num_samples, dset_file_uploads],
            outputs=[dset_status_output, dset_link_output]
        )

    with gr.Tab("3. Entrenamiento"):
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("## ⚙️ Configuración del Entrenamiento")
                training_mode = gr.Dropdown(TRAINING_MODES, label="Modo de Entrenamiento", value=TRAINING_MODES[0])
                with gr.Accordion("📦 Modelo y Repositorio", open=True):
                    model_base_input = gr.Textbox(label="ID del Modelo Base", placeholder="p.ej. 'mistralai/Mistral-7B-v0.1'")
                    tokenizer_name_input = gr.Textbox(label="ID del Tokenizer (opcional)", placeholder="p.ej. si el modelo no tiene tokenizer")
                    repo_name_input = gr.Textbox(label="Nombre del Repositorio de Destino", placeholder="p.ej. 'mi-modelo-afinado'")
                    private_repo = gr.Checkbox(label="Repositorio Privado", value=False)
                    train_from_scratch = gr.Checkbox(label="Entrenar desde Cero", value=False)
                    auto_config_scratch = gr.Checkbox(label="Auto-Configuración", value=True, visible=False)
                    manual_config_scratch = gr.Checkbox(label="Configuración Manual", value=False, visible=False)
                    scratch_architecture = gr.Textbox(label="Arquitectura (p.ej. Llama, Mistral)", value="Llama", visible=False)
                    scratch_vocab_size = gr.Number(label="Tamaño de Vocabulario", value=32000, visible=False)
                    scratch_hidden_size = gr.Number(label="Tamaño Oculto", value=1024, visible=False)
                    scratch_intermediate_size = gr.Number(label="Tamaño Intermedio", value=2048, visible=False)
                    scratch_layers = gr.Number(label="Número de Capas", value=8, visible=False)
                    scratch_heads = gr.Number(label="Cabezas de Atención", value=8, visible=False)
                    scratch_kv_heads = gr.Number(label="Cabezas KV", value=8, visible=False)
                    scratch_block_size = gr.Number(label="Tamaño de Bloque", value=1024, visible=False)
                    scratch_tie_word_embeddings = gr.Checkbox(label="Enlazar Embeddings de Palabras", value=False, visible=False)
                    steps_per_epoch_estimate = gr.Number(label="Estimación de Pasos por Época (para auto-config)", value=1000, visible=False)
                    attention_dropout = gr.Slider(0.0, 0.5, 0.0, label="Dropout de Atención", visible=False)
                    hidden_dropout = gr.Slider(0.0, 0.5, 0.0, label="Dropout Oculto", visible=False)
                    with gr.Accordion("🔄 Fusión de Múltiples Adaptadores (Avanzado)", open=False) as multi_adapter_accordion:
                        enable_multi_adapter_merge = gr.Checkbox(label="Habilitar Fusión Múltiple", value=False)
                        multi_adapter_model_ids = gr.Textbox(label="IDs de Adaptadores (csv)", placeholder="org/adapter1,org/adapter2")
                        multi_adapter_weights = gr.Textbox(label="Pesos (csv)", placeholder="0.5,0.5")
                        multi_adapter_combination_type = gr.Dropdown(["slerp", "linear", "cat", "svd", "dare_linear", "dare_ties", "ties"], label="Tipo de Combinación", value="slerp")
                with gr.Accordion("📚 Dataset", open=True):
                    datasets_hf_text = gr.Textbox(label="Datasets de Hugging Face (csv)", placeholder="p.ej. 'databricks/dolly-15k'")
                    uploads = gr.File(label="Subir Archivos Locales (.jsonl, .csv, .txt)", file_count="multiple")
                    dataset_weights = gr.Textbox(label="Pesos de los Datasets (csv)", placeholder="p.ej. 0.7, 0.3")
                    eval_dataset_hf = gr.Textbox(label="Dataset de Evaluación (opcional)", placeholder="p.ej. 'nombre/dataset_eval'")
                    preview_data_button = gr.Button("Previsualizar Datos Procesados")
                    data_preview_output = gr.Textbox(label="Vista Previa de Datos", lines=8, interactive=False)
                with gr.Accordion("🎓 Hiperparámetros", open=False):
                    with gr.Row():
                        learning_rate = gr.Textbox(label="Tasa de Aprendizaje", value="2e-5")
                        batch_size = gr.Textbox(label="Tamaño de Lote", value="1")
                        gradient_accumulation = gr.Textbox(label="Acumulación de Gradiente", value="8")
                    with gr.Row():
                        block_size = gr.Textbox(label="Longitud de Secuencia", value="1024")
                        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")
                    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")
                         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")
                            adam_beta2 = gr.Textbox(label="Adam Beta2", value="0.999")
                            adam_epsilon = gr.Textbox(label="Adam Epsilon", value="1e-8")
                         group_by_length = gr.Checkbox(label="Agrupar por Longitud", value=False)
                         neftune_noise_alpha = gr.Textbox(label="NEFTune Ruido Alfa (0 para desactivar)", value="0")
                         optim_args = gr.Textbox(label="Argumentos del Optimizador (formato dict)", placeholder="ej: betas=(0.9,0.995)")
                with gr.Accordion("🦋 PEFT (LoRA)", open=True) as peft_accordion:
                    peft = gr.Checkbox(label="Habilitar PEFT/LoRA", value=True)
                    with gr.Row():
                        lora_r = gr.Textbox(label="LoRA r", value="16")
                        lora_alpha = gr.Textbox(label="LoRA alpha", value="32")
                        lora_dropout = gr.Textbox(label="LoRA dropout", value="0.05")
                    auto_find_target_modules = gr.Checkbox(label="Auto-encontrar Módulos de Destino", value=True)
                    target_modules = gr.Textbox(label="Módulos de Destino (csv)", placeholder="q_proj,v_proj", visible=False)
                    modules_to_save = gr.Textbox(label="Módulos a Guardar (csv)", placeholder="embed_tokens,lm_head")
                    with gr.Row():
                        use_dora = gr.Checkbox(label="Usar DoRA", value=False)
                        use_rslora = gr.Checkbox(label="Usar RSLora", value=False)
                    init_lora_weights = gr.Dropdown(["gaussian", "loftq", "pissa"], label="Inicialización de Pesos LoRA", value="gaussian")
                with gr.Accordion("🧹 Procesamiento y Aumentación de Datos", open=False):
                    with gr.Tab("Limpieza y Normalización"):
                        remove_html_tags = gr.Checkbox(label="Eliminar Etiquetas HTML", value=True)
                        normalize_whitespace = gr.Checkbox(label="Normalizar Espacios en Blanco", value=True)
                        remove_urls_emails = gr.Checkbox(label="Eliminar URLs/Emails", value=True)
                        redact_pii = gr.Checkbox(label="Redactar PII", value=True)
                    with gr.Tab("Filtrado"):
                        enable_quality_filter = gr.Checkbox(label="Habilitar Filtros de Calidad", value=True)
                        min_len_input = gr.Slider(1, 100, 10, label="Longitud Mínima (palabras)")
                        max_len_input = gr.Slider(100, 5000, 2000, label="Longitud Máxima (palabras)")
                        rep_threshold_input = gr.Slider(0, 1, 0.2, label="Umbral de Repetición")
                        exclude_keywords_input = gr.Textbox(label="Palabras Clave a Excluir (csv)")
                        bias_keywords_input = gr.Textbox(label="Palabras Clave de Sesgo (csv)", placeholder="p.ej. discriminación,prejuicio")
                        enable_language_filter = gr.Checkbox(label="Habilitar Filtro de Idioma", value=False)
                        allowed_languages = gr.Textbox(label="Idiomas Permitidos (csv)", value="es,en", placeholder="es,en")
                        language_detection_threshold = gr.Slider(0.5, 1.0, 0.95, label="Umbral de Detección de Idioma")
                        enable_toxicity_filter = gr.Checkbox(label="Habilitar Filtro de Toxicidad", value=False)
                        toxicity_threshold = gr.Slider(0.5, 1.0, 0.8, label="Umbral de Toxicidad")
                        enable_coherence_filter = gr.Checkbox(label="Habilitar Filtro de Coherencia (Anti-Gibberish)", value=True)
                        coherence_char_repetition_threshold = gr.Slider(0.1, 0.8, 0.4, label="Umbral de Repetición de Caracteres", info="Máximo ratio de caracteres repetidos permitido")
                        coherence_ngram_repetition_threshold = gr.Slider(0.1, 0.8, 0.3, label="Umbral de Repetición de N-gramas", info="Máximo ratio de patrones repetidos permitido")
                        coherence_entropy_threshold = gr.Slider(0.1, 0.9, 0.5, label="Umbral de Entropía", info="Mínima entropía normalizada requerida")
                        enable_readability_filter = gr.Checkbox(label="Habilitar Filtro de Legibilidad", value=False)
                        min_readability = gr.Slider(0, 100, 30, label="Legibilidad Mínima (Flesch)")
                        max_readability = gr.Slider(0, 100, 100, label="Legibilidad Máxima (Flesch)")
                        enable_stopword_filter = gr.Checkbox(label="Habilitar Filtro de Palabras Vacías", value=False)
                        max_stopword_ratio = gr.Slider(0.0, 1.0, 0.5, label="Ratio Máxima de Palabras Vacías")
                        enable_uniqueness_filter = gr.Checkbox(label="Habilitar Filtro de Unicidad", value=False)
                        min_uniqueness_ratio = gr.Slider(0.0, 1.0, 0.3, label="Ratio Mínima de Unicidad")
                    with gr.Tab("Deduplicación"):
                        deduplication_method = gr.Radio(["Ninguna", "Exacta", "Semántica (MinHash)"], label="Método de Deduplicación", value="Ninguna")
                        minhash_threshold = gr.Slider(0.7, 0.99, 0.85, label="Umbral MinHash")
                        minhash_num_perm = gr.Slider(64, 256, 128, step=16, label="Permutaciones MinHash")
                    with gr.Tab("Aumentación"):
                        enable_back_translation = gr.Checkbox(label="Habilitar Retrotraducción", value=False)
                        bt_model_id = gr.Textbox(label="Modelo de Traducción", value="Helsinki-NLP/opus-mt-en-de")
                        bt_reverse_model_id = gr.Textbox(label="Modelo Inverso", value="Helsinki-NLP/opus-mt-de-en")
                        bt_augmentation_ratio = gr.Slider(0.0, 1.0, 0.1, label="Ratio de Aumentación BT")
                    with gr.Tab("Generación Sintética"):
                        enable_synthetic_data = gr.Checkbox(label="Habilitar Datos Sintéticos", value=False)
                        synthetic_model_id = gr.Textbox(label="ID del Modelo Generador", placeholder="p.ej. 'mistralai/Mistral-7B-Instruct-v0.2'")
                        num_synthetic_samples = gr.Number(label="Número de Muestras", value=1000)
                        synthetic_prompt_template = gr.Textbox(label="Plantilla de Prompt", value="Genera un nuevo ejemplo basado en: {{example_text}}\n\nNuevo ejemplo:", lines=3)
                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:
                        gr.Markdown("Opciones para Text-to-Image aparecerán aquí.")
                with gr.Accordion("📊 Evaluación y Mitigación de Sesgos", open=False):
                    run_evaluation = gr.Checkbox(label="Ejecutar Evaluación", value=False)
                    metric_for_best_model = gr.Textbox(label="Métrica para Mejor Modelo", value="loss", placeholder="loss, accuracy, f1")
                    greater_is_better = gr.Checkbox(label="Mayor es Mejor", 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"]]')
                with gr.Accordion("🔌 Integraciones", open=False):
                    hub_strategy = gr.Dropdown(["every_save", "end", "checkpoint", "all_checkpoints"], label="Estrategia de Subida al Hub", value="every_save")
                    wandb_api_key_input = gr.Textbox(label="Clave API de W&B", type="password")
                    wandb_project_input = gr.Textbox(label="Proyecto W&B")
            with gr.Column(scale=3):
                gr.Markdown("## 📈 Progreso y Resultados")
                with gr.Row():
                    start_training_button = gr.Button("Iniciar Entrenamiento", variant="primary", scale=3)
                    stop_training_button = gr.Button("Detener", variant="stop", visible=False, scale=1)
                training_phase = gr.Label(label="Fase Actual", value="En espera")
                repo_link_output = gr.Markdown(label="Enlace al Repositorio del Modelo")
                final_eval_results = gr.JSON(label="Resultados de Evaluación Final")
                training_logs = gr.Textbox(label="Registros de Entrenamiento", lines=35, interactive=False)
        all_input_components_dict = {
            "training_mode": training_mode, "model_base_input": model_base_input, "tokenizer_name_input": tokenizer_name_input,
            "repo_name_input": repo_name_input, "private_repo": private_repo, "train_from_scratch": train_from_scratch, 
            "auto_config_scratch": auto_config_scratch, "manual_config_scratch": manual_config_scratch,
            "scratch_architecture": scratch_architecture, "scratch_vocab_size": scratch_vocab_size,
            "scratch_hidden_size": scratch_hidden_size, "scratch_intermediate_size": scratch_intermediate_size,
            "scratch_layers": scratch_layers, "scratch_heads": scratch_heads, "scratch_kv_heads": scratch_kv_heads,
            "scratch_block_size": scratch_block_size, "scratch_tie_word_embeddings": scratch_tie_word_embeddings,
            "steps_per_epoch_estimate": steps_per_epoch_estimate, "attention_dropout": attention_dropout, 
            "hidden_dropout": hidden_dropout, "enable_multi_adapter_merge": enable_multi_adapter_merge,
            "multi_adapter_model_ids": multi_adapter_model_ids, "multi_adapter_weights": multi_adapter_weights,
            "multi_adapter_combination_type": multi_adapter_combination_type, "datasets_hf_text": datasets_hf_text,
            "uploads": uploads, "dataset_weights": dataset_weights, "eval_dataset_hf": eval_dataset_hf,
            "learning_rate": learning_rate, "max_steps": max_steps, "batch_size": batch_size, "gradient_accumulation": gradient_accumulation,
            "block_size": block_size, "optimizer": optimizer, "scheduler": scheduler,
            "warmup_ratio": warmup_ratio, "weight_decay": weight_decay, "max_grad_norm": max_grad_norm,
            "logging_steps": logging_steps, "save_steps": save_steps, "save_total_limit": save_total_limit, "resume_from_checkpoint": resume_from_checkpoint,
            "adam_beta1": adam_beta1, "adam_beta2": adam_beta2, "adam_epsilon": adam_epsilon,
            "group_by_length": group_by_length,
            "neftune_noise_alpha": neftune_noise_alpha, "optim_args": optim_args,
            "early_stopping_patience": early_stopping_patience,
            "peft": peft, "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,
            "remove_html_tags": remove_html_tags, "normalize_whitespace": normalize_whitespace, "remove_urls_emails": remove_urls_emails,
            "redact_pii": redact_pii, "enable_quality_filter": enable_quality_filter, "min_len_input": min_len_input,
            "max_len_input": max_len_input, "rep_threshold_input": rep_threshold_input, "exclude_keywords_input": exclude_keywords_input,
            "bias_keywords_input": bias_keywords_input, "enable_language_filter": enable_language_filter,
            "allowed_languages": allowed_languages, "language_detection_threshold": language_detection_threshold,
            "enable_toxicity_filter": enable_toxicity_filter, "toxicity_threshold": toxicity_threshold,
            "enable_coherence_filter": enable_coherence_filter, "coherence_char_repetition_threshold": coherence_char_repetition_threshold,
            "coherence_ngram_repetition_threshold": coherence_ngram_repetition_threshold, "coherence_entropy_threshold": coherence_entropy_threshold,
            "enable_readability_filter": enable_readability_filter, "min_readability": min_readability, "max_readability": max_readability,
            "enable_stopword_filter": enable_stopword_filter, "max_stopword_ratio": max_stopword_ratio,
            "enable_uniqueness_filter": enable_uniqueness_filter, "min_uniqueness_ratio": min_uniqueness_ratio,
            "deduplication_method": deduplication_method, "minhash_threshold": minhash_threshold, "minhash_num_perm": minhash_num_perm,
            "enable_cda": enable_cda, "cda_json_config": cda_json_config,
            "enable_back_translation": enable_back_translation, "bt_model_id": bt_model_id,
            "bt_reverse_model_id": bt_reverse_model_id, "bt_augmentation_ratio": bt_augmentation_ratio,
            "enable_synthetic_data": enable_synthetic_data,
            "synthetic_model_id": synthetic_model_id, "num_synthetic_samples": num_synthetic_samples,
            "synthetic_prompt_template": synthetic_prompt_template,
            "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,
            "run_evaluation": run_evaluation, "metric_for_best_model": metric_for_best_model, 
            "greater_is_better": greater_is_better, "run_perplexity_evaluation": run_perplexity_evaluation,
            "enable_loss_reweighting": enable_loss_reweighting, "reweighting_terms": reweighting_terms, "reweighting_factor": reweighting_factor,
            "hub_strategy": hub_strategy, "wandb_api_key_input": wandb_api_key_input, "wandb_project_input": wandb_project_input,
        }
        all_input_components_list = list(all_input_components_dict.values())
        all_output_components = [training_logs, training_phase, repo_link_output, final_eval_results, start_training_button, stop_training_button]
        preview_data_button.click(
            gradio_preview_data_wrapper,
            inputs=all_input_components_list,
            outputs=[data_preview_output]
        )
        train_from_scratch.change(
            toggle_training_mode_ui,
            inputs=[train_from_scratch],
            outputs=[model_base_input, tokenizer_name_input, multi_adapter_accordion, peft_accordion, 
                     auto_config_scratch, scratch_architecture, manual_config_scratch, scratch_vocab_size,
                     scratch_hidden_size, scratch_intermediate_size, scratch_layers, scratch_heads,
                     scratch_kv_heads, scratch_block_size, scratch_tie_word_embeddings,
                     steps_per_epoch_estimate, attention_dropout, hidden_dropout]
        )
        training_mode.change(
            toggle_task_specific_ui,
            inputs=[training_mode],
            outputs=[classification_labels_ui, dpo_ui, sft_ui, diffusion_ui, peft_accordion]
        )
        sft_format_style.change(
            toggle_sft_format_ui,
            inputs=[sft_format_style],
            outputs=[sft_tool_ui]
        )
        auto_find_target_modules.change(
            toggle_auto_modules_ui,
            inputs=[auto_find_target_modules],
            outputs=[target_modules]
        )
        train_event = start_training_button.click(
            gradio_train_wrapper,
            inputs=all_input_components_list,
            outputs=all_output_components
        )
        stop_training_button.click(fn=None, inputs=None, outputs=None, cancels=[train_event])

    with gr.Tab("4. Inferencia"):
        gr.Markdown("## 🧪 Probar un Modelo del Hub")
        with gr.Row():
            inf_task_mode = gr.Dropdown(TRAINING_MODES, label="Tipo de Tarea", value=TRAINING_MODES[0])
            inf_model_id = gr.Textbox(label="ID del Modelo en el Hub", placeholder="TuUsuario/TuModeloEntrenado")
        with gr.Group():
            with gr.Row():
                with gr.Column(scale=2):
                    inf_text_in = gr.Textbox(label="Entrada de Texto / Prompt", lines=5)
                    inf_context_in = gr.Textbox(label="Contexto (para QA)", lines=3, visible=False)
                    inf_image_in = gr.Image(label="Entrada de Imagen", type="pil", visible=False)
                    inf_audio_in = gr.Audio(label="Entrada de Audio", type="filepath", visible=False)
                    with gr.Accordion("Opciones Avanzadas de Generación", open=False, visible=True) as inf_advanced_options:
                        inf_temperature = gr.Slider(0.1, 2.0, 0.7, label="Temperatura")
                        inf_top_p = gr.Slider(0.1, 1.0, 0.95, label="Top-p")
                        inf_max_new_tokens = gr.Slider(10, 1024, 100, step=1, label="Máximos Tokens Nuevos")
                    with gr.Row():
                        run_inference_btn = gr.Button("Ejecutar Inferencia", variant="primary")
                with gr.Column(scale=3):
                    inf_text_out = gr.Textbox(label="Salida de Texto", lines=15, interactive=False)
        inf_task_mode.change(
            update_inference_ui,
            inputs=[inf_task_mode],
            outputs=[inf_text_in, inf_context_in, inf_image_in, inf_audio_in, inf_advanced_options]
        )
        run_inference_btn.click(
            run_inference,
            inputs=[inf_task_mode, inf_model_id, inf_text_in, inf_context_in, inf_image_in, inf_audio_in, inf_temperature, inf_top_p, inf_max_new_tokens],
            outputs=[inf_text_out, inf_model_id, inf_text_in, inf_context_in, inf_image_in, inf_audio_in]
        )
    with gr.Tab("5. Explicación del Código y Mecanismos Avanzados"):
        gr.Markdown("""
### 🧠 Explicación del Código y Mecanismos Avanzados
""")
        gr.Markdown("#### 1. CORE MECHANISMS")
        gr.Markdown("""
*   PEFT/LoRA: Parameter-Efficient Fine-Tuning. Only low-rank matrices ($A$ and $B$) are trained for low-rank updates ($W' = W + B A$). This drastically reduces trainable parameters.
*   Accelerator: Manages device placement and gradient accumulation for stable large-batch training simulation.
*   Early Stopping: Halts training if validation loss doesn't improve over a set number of steps (`early_stopping_patience`).
*   Gradient Accumulation: Simulates larger batch sizes by accumulating gradients over several forward/backward passes before an optimization step.
*   Gradient Clipping: Limits the maximum norm of the gradients (`max_grad_norm`) to prevent exploding gradients during training.
""")
        gr.Markdown("#### 2. DATA PROCESSING & AUGMENTATION")
        gr.Markdown("""
*   Streaming Datasets: Uses `datasets` streaming mode to handle very large datasets without loading all into RAM.
*   Data Cleaning: Removes HTML tags, normalizes whitespace, redacts PII, and removes URLs/emails.
*   Advanced Filtering: Includes optional filters for text length, word repetition, language detection, and basic toxicity detection (via `unitary/toxic-bert`).
*   Data Augmentation: Supports **Back-Translation (BT)** for introducing paraphrasing variations and **Counterfactual Data Augmentation (CDA)** for controlled bias testing (e.g., swapping gendered pronouns).
*   Synthetic Data Generation: Uses a specified LLM to generate new training examples based on an initial prompt template.
*   Deduplication: Implements both **Exact** and **Semantic (MinHash LSH)** deduplication to prevent data contamination during iterative fine-tuning.
""")
        gr.Markdown("#### 3. TRAINING MODES")
        gr.Markdown("""
*   SFT (Supervised Fine-Tuning): Standard fine-tuning, supports **Conversation** and **Reasoning/Tool Use (CoT)** formatting styles.
*   DPO (Direct Preference Optimization): Trains directly on preference pairs (chosen vs. rejected), using the `trl` library.
*   Task-Specific Heads: Supports **Sequence Classification**, **Token Classification (NER)**, and **Question Answering** by loading appropriate model heads (`AutoModelFor...`).
*   Seq2Seq: For translation/summarization tasks, using `Seq2SeqTrainer`.
""")
        gr.Markdown("#### 4. MODEL INITIALIZATION & ADVANCED TECHNIQUES")
        gr.Markdown("""
*   Model From Scratch: Allows initializing a model (e.g., Llama, Mistral) from a config rather than a pre-trained checkpoint, with optional auto-configuration based on expected training scale.
*   Manual Model Configuration: When training from scratch, users can manually specify low-level configuration parameters (e.g., `vocab_size`, `hidden_size`, `num_hidden_layers`) instead of relying on the automatic scaling based on training steps.
*   Multi-Adapter Merging: Advanced feature to combine multiple existing LoRA adapters into a single, new adapter using weighted averaging (`slerp`, `linear`, etc.).
*   DoRA (Weight-Decomposed Low-Rank Adaptation): A more advanced version of LoRA that can lead to better performance.
*   RSLora (Rank-Stabilized LoRA): A variant of LoRA that adjusts the learning rate based on the rank, improving stability.
*   NEFTune: Adds noise to the embedding layer during training, which can improve the performance of the fine-tuned model.
""")
        gr.Markdown("#### 5. OUTPUT & DEPLOYMENT")
        gr.Markdown("""
*   Hugging Face Hub Integration: All trained artifacts (full model/LoRA adapter) are automatically pushed to a specified repository on the HF Hub using the provided token.
*   Model Card Generation: Automatically generates a `README.md` detailing training parameters and model provenance.
*   Inference Tab: A separate UI for easily testing the trained model with various inputs and generation parameters.
""")

if __name__ == "__main__":
    demo.queue().launch(debug=True, share=True)