Update app.py

app.py

@@ -3,35 +3,34 @@ import subprocess
 import signal
 import tempfile
 from pathlib import Path
-from textwrap import dedent
 from dataclasses import dataclass, field
+from textwrap import dedent
 
 import gradio as gr
 from datasets import load_dataset
-from huggingface_hub import HfApi, ModelCard, whoami
+from huggingface_hub import HfApi
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
 
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# ----------------- Configuración -----------------
+DATASET_NAME = "openerotica/erotiquant3"
 
 @dataclass
 class QuantizationConfig:
     method: str
-    use_imatrix: bool = False
-    imatrix_method: str = "IQ4_NL"
-    train_data: str = ""
+    train_data: str = DATASET_NAME
     quant_embedding: bool = False
-    embedding_tensor_method: str = "Q8_0"
     leave_output: bool = False
     quant_output: bool = False
-    output_tensor_method: str = "Q8_0"
     fp16_model: str = field(default="", init=False)
     quantized_gguf: str = field(default="", init=False)
-    imatrix_file: str = field(default="", init=False)
 
 @dataclass
 class SplitConfig:
     enabled: bool = False
     max_tensors: int = 256
-    max_size: str = None
 
 @dataclass
 class OutputConfig:
@@ -55,76 +54,87 @@ class GGUFConverterError(Exception):
     pass
 
 class HuggingFaceModelProcessor:
-    QUANTIZE_TIMEOUT = 86400
     HF_TO_GGUF_TIMEOUT = 3600
-    IMATRIX_TIMEOUT = 86400
-    SPLIT_TIMEOUT = 3600
+    QUANTIZE_TIMEOUT = 86400
     KILL_TIMEOUT = 5
 
     DOWNLOAD_FOLDER = "./downloads"
     OUTPUT_FOLDER = "./outputs"
-    CALIBRATION_FILE = "calibration_data_v5_rc.txt"
 
     def __init__(self):
         self.HF_TOKEN = os.environ.get("HF_TOKEN")
-        self._create_folder(self.DOWNLOAD_FOLDER)
-        self._create_folder(self.OUTPUT_FOLDER)
-
-    def _create_folder(self, folder_name: str):
-        if not os.path.exists(folder_name):
-            os.makedirs(folder_name)
-        return folder_name
-
-    def _download_dataset(self, dataset_name: str):
-        print(f"Cargando dataset desde HuggingFace Hub: {dataset_name}")
-        dataset = load_dataset(dataset_name, use_auth_token=self.HF_TOKEN)
-        return dataset
-
-    def _download_model(self, processing_config: ModelProcessingConfig):
-        print(f"Descargando modelo {processing_config.model_name}")
-        if os.path.exists(processing_config.quant_config.fp16_model):
-            print("FP16 ya existe, omitiendo conversión.")
-            return processing_config.quant_config.fp16_model
-
-        with tempfile.TemporaryDirectory(dir=self.DOWNLOAD_FOLDER) as tmpdir:
-            local_dir = f"{Path(tmpdir)}/{processing_config.model_name}"
-            api = HfApi(token=processing_config.token)
-            pattern = "*.safetensors"
-            api.snapshot_download(repo_id=processing_config.model_id, local_dir=local_dir, allow_patterns=[pattern])
-            convert_command = [
-                "python3", "/app/convert_hf_to_gguf.py", local_dir,
-                "--outtype", "f16", "--outfile", processing_config.quant_config.fp16_model
-            ]
-            process = subprocess.Popen(convert_command, shell=False, stderr=subprocess.STDOUT)
+        os.makedirs(self.DOWNLOAD_FOLDER, exist_ok=True)
+        os.makedirs(self.OUTPUT_FOLDER, exist_ok=True)
+
+    # ----------------- Entrenamiento -----------------
+    def _train_model(self, model_id: str, dataset_name: str, outdir: str):
+        print(f"Iniciando entrenamiento de {model_id} con dataset {dataset_name}")
+        dataset = load_dataset(dataset_name, split="train", use_auth_token=self.HF_TOKEN)
+
+        tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=self.HF_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=self.HF_TOKEN)
+
+        def tokenize_function(examples):
+            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
+
+        tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)
+
+        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+        train_args = TrainingArguments(
+            output_dir=f"{outdir}/{Path(model_id).name}_trained",
+            overwrite_output_dir=True,
+            num_train_epochs=1,  # ajustar según recursos
+            per_device_train_batch_size=2,
+            save_total_limit=1,
+            logging_dir=f"{outdir}/logs",
+            logging_steps=10,
+            save_steps=200,
+            report_to=[],
+        )
+
+        trainer = Trainer(
+            model=model,
+            args=train_args,
+            train_dataset=tokenized_dataset,
+            data_collator=data_collator,
+        )
+
+        trainer.train()
+        trainer.save_model(train_args.output_dir)
+        tokenizer.save_pretrained(train_args.output_dir)
+
+        print("Entrenamiento finalizado.")
+        return train_args.output_dir
+
+    # ----------------- Conversión a FP16 -----------------
+    def _convert_to_fp16(self, model_dir: str, out_fp16_path: str):
+        print(f"Convirtiendo modelo a FP16: {out_fp16_path}")
+        convert_command = [
+            "python3", "/app/convert_hf_to_gguf.py",
+            model_dir,
+            "--outtype", "f16",
+            "--outfile", out_fp16_path
+        ]
+        process = subprocess.Popen(convert_command, shell=False, stderr=subprocess.STDOUT)
+        try:
+            process.wait(timeout=self.HF_TO_GGUF_TIMEOUT)
+        except subprocess.TimeoutExpired:
+            process.send_signal(signal.SIGINT)
             try:
-                process.wait(timeout=self.HF_TO_GGUF_TIMEOUT)
+                process.wait(timeout=self.KILL_TIMEOUT)
             except subprocess.TimeoutExpired:
-                process.send_signal(signal.SIGINT)
-                try:
-                    process.wait(timeout=self.KILL_TIMEOUT)
-                except subprocess.TimeoutExpired:
-                    process.kill()
-                raise GGUFConverterError("Error convirtiendo a FP16: timeout")
-            if process.returncode != 0:
-                raise GGUFConverterError(f"Error convirtiendo a FP16: code={process.returncode}")
-            print("Modelo convertido a FP16 correctamente")
-            return processing_config.quant_config.fp16_model
-
-    def _quantize_model(self, quant_config: QuantizationConfig):
-        quantize_cmd = ["llama-quantize"]
-        if quant_config.quant_embedding:
-            quantize_cmd.extend(["--token-embedding-type", quant_config.embedding_tensor_method])
-        if quant_config.leave_output:
-            quantize_cmd.append("--leave-output-tensor")
-        else:
-            if quant_config.quant_output:
-                quantize_cmd.extend(["--output-tensor-type", quant_config.output_tensor_method])
-        if quant_config.use_imatrix:
-            raise NotImplementedError("imatrix no implementado para esta demo automática")
-        quantize_cmd.append(quant_config.fp16_model)
-        quantize_cmd.append(quant_config.quantized_gguf)
-        quantize_cmd.append(quant_config.method)
-
+                process.kill()
+            raise GGUFConverterError("Error convirtiendo a FP16: timeout")
+        if process.returncode != 0:
+            raise GGUFConverterError(f"Error FP16: código {process.returncode}")
+        print("Conversión a FP16 completa.")
+        return out_fp16_path
+
+    # ----------------- Cuantización -----------------
+    def _quantize_model(self, fp16_path: str, gguf_path: str, method: str):
+        print(f"Cuantizando modelo con método {method}")
+        quantize_cmd = ["llama-quantize", fp16_path, gguf_path, method]
         process = subprocess.Popen(quantize_cmd, shell=False, stderr=subprocess.STDOUT)
         try:
             process.wait(timeout=self.QUANTIZE_TIMEOUT)
@@ -137,25 +147,12 @@ class HuggingFaceModelProcessor:
             raise GGUFConverterError("Error cuantizando: timeout")
         if process.returncode != 0:
             raise GGUFConverterError(f"Error cuantizando: code={process.returncode}")
-        print("Cuantización completada")
-        return quant_config.quantized_gguf
-
-    def _create_repo(self, processing_config: ModelProcessingConfig):
-        api = HfApi(token=processing_config.token)
-        new_repo_url = api.create_repo(repo_id=processing_config.output_config.repo_name, exist_ok=True, private=processing_config.output_config.private_repo)
-        processing_config.new_repo_url = new_repo_url.url
-        processing_config.new_repo_id = new_repo_url.repo_id
-        print("Repositorio creado:", processing_config.new_repo_url)
-        return new_repo_url
-
-    def run_full_pipeline(self, token, model_id, model_name, dataset_name):
+        print("Cuantización completada.")
+        return gguf_path
+
+    def run_full_pipeline(self, token, model_id, model_name):
         logs = []
         try:
-            # 1. Cargar dataset
-            dataset = self._download_dataset(dataset_name)
-            logs.append(f"Dataset cargado: {dataset_name}")
-
-            # 2. Configuración inicial
             outdir = self.OUTPUT_FOLDER
             quant_config = QuantizationConfig(method="Q4_0")
             quant_config.fp16_model = f"{outdir}/{model_name}.f16"
@@ -163,24 +160,27 @@ class HuggingFaceModelProcessor:
             split_config = SplitConfig()
             output_config = OutputConfig(private_repo=False, repo_name=f"{model_name}-gguf")
             processing_config = ModelProcessingConfig(
-                token=token, model_id=model_id, model_name=model_name, outdir=outdir,
-                quant_config=quant_config, split_config=split_config, output_config=output_config
+                token=token, model_id=model_id, model_name=model_name,
+                outdir=outdir, quant_config=quant_config, split_config=split_config,
+                output_config=output_config
             )
 
-            # 3. Descargar modelo
-            self._download_model(processing_config)
-            logs.append("Modelo descargado y convertido a FP16")
+            # 1. Entrenar modelo
+            trained_dir = self._train_model(model_id, DATASET_NAME, outdir)
+            logs.append("Entrenamiento completado")
+
+            # 2. Convertir a FP16
+            self._convert_to_fp16(trained_dir, quant_config.fp16_model)
+            logs.append("Conversión a FP16 completada")
 
-            # 4. Cuantizar modelo
-            self._quantize_model(quant_config)
+            # 3. Cuantizar
+            self._quantize_model(quant_config.fp16_model, quant_config.quantized_gguf, quant_config.method)
             logs.append("Modelo cuantizado a GGUF")
 
-            # 5. Crear repo
-            self._create_repo(processing_config)
-            logs.append(f"Repositorio creado: {processing_config.new_repo_url}")
+            logs.append("✅ Pipeline completada correctamente")
 
         except Exception as e:
-            logs.append(f"ERROR: {e}")
+            logs.append(f"❌ ERROR: {e}")
 
         return "\n".join(logs)
 
@@ -188,16 +188,15 @@ class HuggingFaceModelProcessor:
 processor = HuggingFaceModelProcessor()
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Pipeline Automática GGUF desde HuggingFace Hub")
-    dataset_input = gr.Textbox(label="Nombre del dataset HuggingFace", placeholder="openerotica/erotiquant3")
+    gr.Markdown("## Pipeline Automática GGUF con entrenamiento (dataset openerotica/erotiquant3)")
     model_input = gr.Textbox(label="ID del modelo HF", placeholder="ochoa/your-model")
-    token_input = gr.Textbox(label="Tu token HF (opcional, si está en HF_TOKEN puede dejarse vacío)", type="password")
-    run_button = gr.Button("Ejecutar pipeline automática")
+    token_input = gr.Textbox(label="Tu token HF (opcional si ya está en HF_TOKEN)", type="password")
+    run_button = gr.Button("Entrenar y cuantizar automáticamente")
     output_logs = gr.Textbox(label="Logs", lines=20)
 
     run_button.click(
-        fn=lambda token, model_id, model_name, dataset_name: processor.run_full_pipeline(token, model_id, model_name, dataset_name),
-        inputs=[token_input, model_input, model_input, dataset_input],
+        fn=lambda token, model_id: processor.run_full_pipeline(token, model_id, Path(model_id).name),
+        inputs=[token_input, model_input],
         outputs=[output_logs]
     )