File size: 8,378 Bytes
bbe1432
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
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