Update app.py

app.py

@@ -7,7 +7,9 @@ from dataclasses import dataclass, field
 import gradio as gr
 from datasets import load_dataset
 from huggingface_hub import HfApi, Repository
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
+from trl import PPOTrainer, PPOConfig
+from trl import AutoModelForCausalLMWithValueHead
 
 os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -59,46 +61,44 @@ class HuggingFaceModelProcessor:
         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, token: str):
-        print(f"Iniciando entrenamiento de {model_id} con dataset {dataset_name}")
+    # ----------------- Entrenamiento con PPO -----------------
+    def _train_model_with_ppo(self, model_id: str, dataset_name: str, outdir: str, token: str):
+        print(f"Iniciando entrenamiento con PPO de {model_id} con dataset {dataset_name}")
         dataset = load_dataset(dataset_name, split="train", token=token)
 
         tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
-        model = AutoModelForCausalLM.from_pretrained(model_id, token=token)
+        model = AutoModelForCausalLMWithValueHead.from_pretrained(model_id, token=token)
 
+        # Tokenización
         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,
+        # Configuración de PPO
+        ppo_config = PPOConfig(
+            model_name=model_id,
+            train_batch_size=2,
+            gradient_accumulation_steps=1,
+            learning_rate=5e-6,
+            max_length=512,
             num_train_epochs=1,
-            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(
+        ppo_trainer = PPOTrainer(
             model=model,
-            args=train_args,
+            args=ppo_config,
             train_dataset=tokenized_dataset,
-            data_collator=data_collator,
+            tokenizer=tokenizer,
         )
 
-        trainer.train()
-        trainer.save_model(train_args.output_dir)
-        tokenizer.save_pretrained(train_args.output_dir)
+        # Entrenamiento
+        ppo_trainer.train()
+        ppo_trainer.save_model(outdir)
+        tokenizer.save_pretrained(outdir)
 
-        print("Entrenamiento finalizado.")
-        return train_args.output_dir
+        print("Entrenamiento PPO finalizado.")
+        return outdir
 
     # ----------------- Conversión a FP16 -----------------
     def _convert_to_fp16(self, model_dir: str, out_fp16_path: str):
@@ -167,8 +167,8 @@ class HuggingFaceModelProcessor:
                 output_config=output_config
             )
 
-            # 1. Entrenar modelo
-            trained_dir = self._train_model(model_id, DATASET_NAME, outdir, token)
+            # 1. Entrenar modelo con PPO
+            trained_dir = self._train_model_with_ppo(model_id, DATASET_NAME, outdir, token)
             logs.append("Entrenamiento completado")
 
             # 2. Convertir a FP16
@@ -194,7 +194,7 @@ class HuggingFaceModelProcessor:
 processor = HuggingFaceModelProcessor()
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Pipeline Automática GGUF con entrenamiento y publicación HF")
+    gr.Markdown("## Pipeline Automática GGUF con entrenamiento PPO y publicación HF")
     model_input = gr.Textbox(label="ID del modelo HF (para entrenamiento)", placeholder="ochoa/your-model")
     repo_input = gr.Textbox(label="Nombre del repo HF para publicar", placeholder="usuario/nuevo-modelo")
     token_input = gr.Textbox(label="Tu token HF", type="password")