model_utils.py

import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import re
from peft import PeftModel


def get_model_and_tokenizer(model_name="prometheus-eval/prometheus-7b-v2.0" ):
    """
    Carga el modelo Prometheus y su tokenizador asociado desde Hugging Face.
    
    Esta función es esencial para el hackathon ya que inicializa el evaluador LLM-as-a-Judge.
    Recuerda configurar tu token de Hugging Face de antemano.
    
    Args:
        model_name (str): La versión específica del modelo de Prometheus a cargar.
        
    Returns:
        model, tokenizer: Tupla con el modelo y el tokenizador listos para realizar inferencias.
    """
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        print("Warning: HF_TOKEN not found in environment variables.")
        
    print(f"Loading model: {model_name}...")
    
     # 1. Cargar y configurar el Tokenizador
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
    
    # Configuramos el pad_token si no existe (común en Mistral/Llama)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Padding a la izquierda es obligatorio para modelos decodificadores (CausalLM) 
    # cuando se hace inferencia en batches
    tokenizer.padding_side = "left" 
    
    # 2. Cargar el Modelo
    model = AutoModelForCausalLM.from_pretrained(
         model_name,
         token=hf_token,
         device_map="auto",
         # dtype=torch.float16, # Media precisión para ganar velocidad y ahorrar VRAM # No necesario dado que usamos cuantización a 4-bit
         low_cpu_mem_usage=True,
         trust_remote_code=True, # Añadido
         use_safetensors=True, # Añadido
         #safetensors_filename="gptq_model-4bit-128g.safetensors"
    )
    
    
    return model, tokenizer



    
def split_model_reason_result(sample, output_suffix : str = "model", input_col: str = "model_output")->dict:
    """
    Post-procesa la salida del modelo para separar la explicación de la puntuación.
    
    Busca la etiqueta '[RESULT]' para dividir el texto. Si no la encuentra, 
    asume que todo el texto es el razonamiento y devuelve un resultado nulo.

    Args:
        sample (dict | str): Ejemplo que contiene 'model_output'.
        output_suffix (str): Sufijo para nombrar la columna de salida.
        input_col (str): Nombre de la columna de entrada.

    Returns:
        dict: Diccionario con las claves 'reason' (explicación) y 'result' (puntuación limpia).
    """

    output = sample.get(input_col, "") if not isinstance(sample, str) else sample
    
    if "[RESULT]" in output:
        # Dividimos por la última aparición del tag para evitar errores
        parts = output.rsplit("[RESULT]", 1)
        reason = parts[0].strip()
        result_raw = parts[1].strip()
        
        # Limpieza mediante regex para capturar solo el dígito (evita puntos finales, etc.)
        score_match = re.search(r'(\d+)', result_raw)
        result = score_match.group(1) if score_match else result_raw
    else:
        reason = output.strip()
        result = None
    
    return {
        f"{output_suffix}_reason": reason,
        f"{output_suffix}_pred": result
    }



def model_predict(model, tokenizer, prompt, max_new_tokens =200, temperature=0.7):
    """
    Realiza una inferencia simple para un único prompt utilizando el modelo y tokenizador proporcionados.

    Esta función prepara el texto, lo envía al dispositivo donde reside el modelo (GPU/CPU) 
    y genera una respuesta de forma determinista. Es ideal para pruebas rápidas o 
    validaciones unitarias durante la hackathon.

    Args:
        model (transformers.PreTrainedModel): El modelo de lenguaje ya cargado.
        tokenizer (transformers.PreTrainedTokenizer): El tokenizador correspondiente al modelo.
        prompt (str): El texto de entrada o instrucción para el modelo.

    Returns:
        str: El texto generado por el modelo, limpio de tokens especiales y del prompt original.
    """
    # 1. Identificar el dispositivo del modelo (soporta device_map="auto")
    device = model.device 
    
    # 2. Tokenizar y mover tensores al dispositivo correcto
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    # 3. Generación determinista (do_sample=False para evitar variabilidad en pruebas)
    with torch.no_grad():
        outputs = model.generate(
            **inputs, 
            max_new_tokens=max_new_tokens, 
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id else tokenizer.eos_token_id
        )

    # 4. Decodificar solo la parte nueva (ignorando los tokens del prompt)
    input_length = inputs["input_ids"].shape[1]
    prediction = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)

    return prediction.strip()



def model_predict_batched(model, tokenizer, batch, input_col = "user_content", 
                            temperature = 0.1, max_new_tokens = 1000, completion_colname = "model_output"):
    """
    Realiza inferencia en lotes (batches) sobre un conjunto de prompts.
    
    Esta función es más eficiente que `model_predict` cuando se procesan múltiples ejemplos a la vez,
    ya que aprovecha el procesamiento en paralelo de la GPU. Aplica el template de chat 
    del tokenizador automáticamente.
    
    Args:
        model (transformers.PreTrainedModel): El modelo cargado.
        tokenizer (transformers.PreTrainedTokenizer): El tokenizador correspondiente.
        batch (dict o pd.DataFrame): El lote de datos de entrada.
        input_col (str, opcional): El nombre de la columna que contiene los prompts de usuario. Por defecto "user_content".
        temperature (float, opcional): Parámetro de temperatura para controlar la aleatoriedad. Por defecto 0.1.
        max_new_tokens (int, opcional): Límite máximo de tokens a generar. Por defecto 1000.
        completion_colname (str, opcional): Nombre de la columna de salida. Por defecto "model_output".
        
    Returns:
        dict: Diccionario que contiene una lista con las respuestas generadas bajo la clave f"{completion_colname}".
    """
    # 1. Detectamos el dispositivo de entrada (donde está la primera c  apa)
    model_device = model.device 
    
    messages_list = [[{"role": "user", "content": p}] for p in batch[input_col]]


    # 2. IMPORTANTE: Pedimos que devuelva un diccionario completo (return_dict=True)
    inputs = tokenizer.apply_chat_template(
        messages_list,
        add_generation_prompt=True,
        tokenize=True,
        return_tensors="pt",
        padding=True,
        return_dict=True # Esto asegura que tengamos input_ids y attention_mask
    ).to(model_device)

    with torch.no_grad():
        generated_ids = model.generate(
            **inputs, # Ahora inputs es un dict con todo en la GPU correcta
            max_new_tokens=max_new_tokens, 
            do_sample=True,
            temperature=temperature,
            pad_token_id=tokenizer.pad_token_id
        )
    
    input_length = inputs["input_ids"].shape[1]
    decoded_outputs = tokenizer.batch_decode(
        generated_ids[:, input_length:], 
        skip_special_tokens=True
    )
    
    return {f"{completion_colname}": decoded_outputs}




def load_lora_model(model_name, model_path):
    """
    Carga un modelo base y le aplica los pesos ajustados de un entrenamiento LoRA (PEFT).
    
    Durante el hackathon, usarás esta función para cargar tu propio modelo afinao (Fine-Tuned)
    y comparar sus evaluaciones con las del modelo original.
    
    Args:
        model_name (str): Nombre o ruta del modelo base original (p. ej., "prometheus-eval/prometheus-7b-v2.0").
        model_path (str): Ruta donde se encuentran guardados los adaptadores LoRA entrenados.
        
    Returns:
        model, tokenizer: Tupla con el modelo ajustado y su tokenizador.
    """

    # 1. Load the original BASE model (the one you started with)
    base_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
    
    # 2. Load the Tokenizer (now that you've saved it to the FT path)
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    
    # 3. Load the LoRA adapters onto the base model
    model = PeftModel.from_pretrained(base_model, model_path)
    return model, tokenizer