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