Update app.py

app.py

@@ -8,6 +8,14 @@ from PIL import Image
 import gradio as gr
 from huggingface_hub import login, HfApi
 
+GRADIO_LOG = ""
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+device = torch.device("cpu")
+checkpoint_path = pathlib.Path("/tmp/gemma_pytorch_models/checkpoint.pth")
+save_dir = pathlib.Path("/tmp/gemma_pytorch_models/")
+save_dir.mkdir(exist_ok=True, parents=True)
+
 def log_message(msg, level="info"):
     global GRADIO_LOG
     if level == "info":
@@ -18,14 +26,6 @@ def log_message(msg, level="info"):
         logging.debug(msg)
     GRADIO_LOG += msg + "\n"
 
-GRADIO_LOG = ""
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
-os.environ["CUDA_VISIBLE_DEVICES"] = ""
-device = torch.device("cpu")
-checkpoint_path = pathlib.Path("/tmp/gemma_pytorch_models/checkpoint.pth")
-save_dir = pathlib.Path("/tmp/gemma_pytorch_models/")
-save_dir.mkdir(exist_ok=True, parents=True)
-
 def load_checkpoint(model):
     if checkpoint_path.exists():
         try:
@@ -351,94 +351,43 @@ def distillation_loss(student_logits, teacher_logits):
     loss_hard = F.mse_loss(student_logits, teacher_logits)
     return alpha * loss_soft + (1 - alpha) * loss_hard
 
-optimizer = optim.Adam(student_model.parameters(), lr=1e-4)
-student_model.train()
-print("Inicio del entrenamiento por destilación")
-for epoch in range(1):
-    print(f"Época {epoch+1}")
-    for text in tqdm([], desc="Entrenamiento destilación"):
-        optimizer.zero_grad()
-        inputs = student_tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=16)
-        with torch.no_grad():
-            t1_inputs = teacher1_tokenizer(text=[text], return_tensors="pt", padding="max_length", truncation=True, max_length=16)
-            t2_inputs = teacher2_tokenizer(text=[text], return_tensors="pt", padding="max_length", truncation=True, max_length=16)
-            teacher1_out = teacher1_model(**t1_inputs, max_new_tokens=10)
-            teacher2_out = teacher2_model(**t2_inputs, max_new_tokens=10)
-        teacher_logits = (teacher1_out.logits + teacher2_out.logits) / 2.0
-        student_out = student_model(**inputs)
-        student_logits = student_out.logits
-        loss = distillation_loss(student_logits, teacher_logits)
-        loss.backward()
-        optimizer.step()
-        print(f"Texto: {text} | Loss: {loss.item():.6f}")
-print("Entrenamiento completado")
-save_checkpoint(student_model)
-print("Aplicando cuantización dinámica al modelo student")
-student_quantized = torch.quantization.quantize_dynamic(student_model, {nn.Linear}, dtype=torch.qint8)
-def save_model_state(model, filename):
-    state_dict = model.state_dict()
-    np_state = {k: v.cpu().numpy() for k, v in state_dict.items()}
-    save_file(np_state, str(filename))
-for _ in tqdm(range(1), desc="Guardando modelos"):
-    pass
-teacher1_file = save_dir / "gemma_teacher_model_quant.safetensors"
-teacher2_file = save_dir / "llama_teacher_model_quant.safetensors"
-student_file = save_dir / "gemma_student_model_quant.safetensors"
-save_model_state(teacher1_quantized, teacher1_file)
-save_model_state(teacher2_quantized, teacher2_file)
-save_model_state(student_quantized, student_file)
-print(f"Modelos guardados en {save_dir}")
-try:
-    user_info = HfApi().whoami(token=hf_token_input)
-    username = user_info["name"]
-    repo_id_t1 = f"{username}/gemma-teacher-pytorch-safetensors"
-    HfApi().create_repo(repo_id_t1, token=hf_token_input, exist_ok=True)
-    HfApi().upload_file(token=hf_token_input, path_or_fileobj=str(teacher1_file), path_in_repo="gemma_teacher_model_quant.safetensors", repo_id=repo_id_t1)
-    repo_id_t2 = f"{username}/llama-teacher-pytorch-safetensors"
-    HfApi().create_repo(repo_id_t2, token=hf_token_input, exist_ok=True)
-    HfApi().upload_file(token=hf_token_input, path_or_fileobj=str(teacher2_file), path_in_repo="llama_teacher_model_quant.safetensors", repo_id=repo_id_t2)
-    repo_id_student = f"{username}/gemma-student-pytorch-safetensors"
-    HfApi().create_repo(repo_id_student, token=hf_token_input, exist_ok=True)
-    HfApi().upload_file(token=hf_token_input, path_or_fileobj=str(student_file), path_in_repo="gemma_student_model_quant.safetensors", repo_id=repo_id_student)
-    print("Modelos subidos al Hub de Hugging Face")
-except Exception as e:
-    print(f"Error subiendo modelos: {e}")
-
 def run_pipeline(teacher1_id, teacher2_id, student_id, hf_token_input):
+    global GRADIO_LOG
+    GRADIO_LOG = ""
     try:
         login(token=hf_token_input)
         api = HfApi()
         proc = AutoProcessor.from_pretrained(teacher1_id)
-        t1_model = AutoModelForImageTextToText.from_pretrained(teacher1_id)
-        t1_model.to(device); t1_model.eval()
-        t1_tokenizer = proc.tokenizer if hasattr(proc, "tokenizer") else proc
-        t2_tokenizer = AutoTokenizer.from_pretrained(teacher2_id)
-        t2_model = AutoModelForCausalLM.from_pretrained(teacher2_id)
-        t2_model.to(device); t2_model.eval()
-        s_tokenizer = AutoTokenizer.from_pretrained(student_id)
-        if s_tokenizer.pad_token is None:
-            s_tokenizer.add_special_tokens({'pad_token': s_tokenizer.eos_token if s_tokenizer.eos_token is not None else "[PAD]"})
-        s_model = AutoModelForCausalLM.from_pretrained(student_id)
-        s_model.to(device)
+        teacher1_model = AutoModelForImageTextToText.from_pretrained(teacher1_id)
+        teacher1_model.to(device); teacher1_model.eval()
+        teacher1_tokenizer = proc.tokenizer if hasattr(proc, "tokenizer") else proc
+        teacher2_tokenizer = AutoTokenizer.from_pretrained(teacher2_id)
+        teacher2_model = AutoModelForCausalLM.from_pretrained(teacher2_id)
+        teacher2_model.to(device); teacher2_model.eval()
+        student_tokenizer = AutoTokenizer.from_pretrained(student_id)
+        if student_tokenizer.pad_token is None:
+            student_tokenizer.add_special_tokens({'pad_token': student_tokenizer.eos_token if student_tokenizer.eos_token is not None else "[PAD]"})
+        student_model = AutoModelForCausalLM.from_pretrained(student_id)
+        student_model.to(device)
         print("Modelos cargados")
     except Exception as e:
         return f"Error cargando modelos: {e}"
     try:
-        t1_task = getattr(t1_model.config, "task_type", "image-to-text")
+        t1_task = getattr(teacher1_model.config, "task_type", "image-to-text")
     except Exception:
         t1_task = "image-to-text"
     try:
-        t2_task = getattr(t2_model.config, "task", "text-generation")
+        t2_task = getattr(teacher2_model.config, "task", "text-generation")
     except Exception:
         t2_task = "text-generation"
     try:
-        s_task = getattr(s_model.config, "task", "text-generation")
+        s_task = getattr(student_model.config, "task", "text-generation")
     except Exception:
         s_task = "text-generation"
     try:
-        pipe_t1 = pipeline(t1_task, model=t1_model, tokenizer=t1_tokenizer, device=-1)
-        pipe_t2 = pipeline(t2_task, model=t2_model, tokenizer=t2_tokenizer, device=-1)
-        pipe_s = pipeline(s_task, model=s_model, tokenizer=s_tokenizer, device=-1)
+        pipe_t1 = pipeline(t1_task, model=teacher1_model, tokenizer=teacher1_tokenizer, device=-1)
+        pipe_t2 = pipeline(t2_task, model=teacher2_model, tokenizer=teacher2_tokenizer, device=-1)
+        pipe_s = pipeline(s_task, model=student_model, tokenizer=student_tokenizer, device=-1)
         print("Pipelines creados:")
         print("Teacher1:", pipe_t1.task)
         print("Teacher2:", pipe_t2.task)
@@ -455,6 +404,60 @@ def run_pipeline(teacher1_id, teacher2_id, student_id, hf_token_input):
         print("Student:", out_s)
     except Exception as e:
         return f"Error en ejecución de pipelines de prueba: {e}"
+
+    optimizer = optim.Adam(student_model.parameters(), lr=1e-4)
+    student_model.train()
+    print("Inicio del entrenamiento por destilación")
+    texts = ["Texto de ejemplo 1", "Texto de ejemplo 2"]
+    for epoch in range(1):
+        print(f"Época {epoch+1}")
+        for text in tqdm(texts, desc="Entrenamiento destilación"):
+            optimizer.zero_grad()
+            inputs = student_tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=16)
+            with torch.no_grad():
+                t1_inputs = teacher1_tokenizer(text=[text], return_tensors="pt", padding="max_length", truncation=True, max_length=16)
+                t2_inputs = teacher2_tokenizer(text=[text], return_tensors="pt", padding="max_length", truncation=True, max_length=16)
+                teacher1_out = teacher1_model(**t1_inputs, max_new_tokens=10)
+                teacher2_out = teacher2_model(**t2_inputs, max_new_tokens=10)
+                teacher_logits = (teacher1_out.logits + teacher2_out.logits) / 2.0
+            student_out = student_model(**inputs)
+            student_logits = student_out.logits
+            loss = distillation_loss(student_logits, teacher_logits)
+            loss.backward()
+            optimizer.step()
+            print(f"Texto: {text} | Loss: {loss.item():.6f}")
+    print("Entrenamiento completado")
+    save_checkpoint(student_model)
+    print("Aplicando cuantización dinámica al modelo student")
+    student_quantized = torch.quantization.quantize_dynamic(student_model, {nn.Linear}, dtype=torch.qint8)
+    def save_model_state(model, filename):
+        state_dict = model.state_dict()
+        np_state = {k: v.cpu().numpy() for k, v in state_dict.items()}
+        save_file(np_state, str(filename))
+    for _ in tqdm(range(1), desc="Guardando modelos"):
+        pass
+    teacher1_file = save_dir / "gemma_teacher_model_quant.safetensors"
+    teacher2_file = save_dir / "llama_teacher_model_quant.safetensors"
+    student_file = save_dir / "gemma_student_model_quant.safetensors"
+    save_model_state(teacher1_model, teacher1_file)
+    save_model_state(teacher2_model, teacher2_file)
+    save_model_state(student_quantized, student_file)
+    print(f"Modelos guardados en {save_dir}")
+    try:
+        user_info = HfApi().whoami(token=hf_token_input)
+        username = user_info["name"]
+        repo_id_t1 = f"{username}/gemma-teacher-pytorch-safetensors"
+        HfApi().create_repo(repo_id_t1, token=hf_token_input, exist_ok=True)
+        HfApi().upload_file(token=hf_token_input, path_or_fileobj=str(teacher1_file), path_in_repo="gemma_teacher_model_quant.safetensors", repo_id=repo_id_t1)
+        repo_id_t2 = f"{username}/llama-teacher-pytorch-safetensors"
+        HfApi().create_repo(repo_id_t2, token=hf_token_input, exist_ok=True)
+        HfApi().upload_file(token=hf_token_input, path_or_fileobj=str(teacher2_file), path_in_repo="llama_teacher_model_quant.safetensors", repo_id=repo_id_t2)
+        repo_id_student = f"{username}/gemma-student-pytorch-safetensors"
+        HfApi().create_repo(repo_id_student, token=hf_token_input, exist_ok=True)
+        HfApi().upload_file(token=hf_token_input, path_or_fileobj=str(student_file), path_in_repo="gemma_student_model_quant.safetensors", repo_id=repo_id_student)
+        print("Modelos subidos al Hub de Hugging Face")
+    except Exception as e:
+        print(f"Error subiendo modelos: {e}")
     return GRADIO_LOG + "\nProceso completado."
 
 iface = gr.Interface(
@@ -470,4 +473,4 @@ iface = gr.Interface(
     description="Ingrese los IDs de los modelos y su token HF para ejecutar el pipeline automáticamente."
 )
 
-iface.launch()
+iface.launch()