Update app.py

app.py

@@ -7,64 +7,42 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from huggingface_hub import login
 import torch
 import json
+import os
 from datetime import datetime
 
+# Configurar variable de entorno para evitar la fragmentación de memoria en CUDA
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
 # Autenticación con Fireworks en Hugging Face
 huggingface_token = st.secrets["FIREWORKS"]
 login(huggingface_token)
 
-# Configurar modelo Fireworks con cuantización int8
+# Configurar modelo Fireworks con cuantización int8 y offload en la CPU
 quant_config = BitsAndBytesConfig(
     load_in_8bit=True,  # Activar la carga en int8
-    llm_int8_enable_fp32_cpu_offload=False,  # # Desactivar el offload en la CPU
+    llm_int8_enable_fp32_cpu_offload=True,  # Permitir offload en la CPU
     quantization_scheme="gptq"  # Especificar el esquema GPTQ
 )
 
 model_id = "fireworks-ai/firefunction-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
-    model_id, 
-    device_map={"": 0}, # Forzar la carga del modelo en la GPU 0
+    model_id,
+    device_map="auto",  # Permitir el offload automático entre CPU y GPU
     torch_dtype=torch.float16,
     quantization_config=quant_config
 )
 
-# Asegurarte de mover el modelo a CUDA
-model.to("cuda")
-
 # Establecer el token de relleno
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
-# Definir funciones específicas para Fireworks
-function_spec = [
-    {
-        "name": "calculate_cosine_similarity",
-        "description": "Calculate the cosine similarity between two strings.",
-        "parameters": {
-            "type": "object",
-            "properties": {
-                "query": {
-                    "type": "string",
-                    "description": "The main query string for similarity calculation"
-                },
-                "job_title": {
-                    "type": "string",
-                    "description": "The job title to compare with the query"
-                }
-            },
-            "required": ["query", "job_title"]
-        }
-    }
-]
-functions = json.dumps(function_spec, indent=4)
-
 # Crear pipeline para generación de texto con Fireworks
 fireworks_pipeline = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_new_tokens=50  # Reducir max_new_tokens para acelerar
+    max_new_tokens=20  # Reducir max_new_tokens para minimizar el uso de memoria
 )
 
 # Adaptar el pipeline a LangChain
@@ -84,7 +62,7 @@ if uploaded_file is not None:
         job_titles = df['job_title'].tolist()
 
         # Procesar en lotes para optimización
-        batch_size = 16  # Ajusta según la memoria de la GPU
+        batch_size = 4  # Reducir batch size para minimizar el uso de memoria
         job_titles_batches = [job_titles[i:i+batch_size] for i in range(0, len(job_titles), batch_size)]
 
         # Definir el prompt para Fireworks
@@ -117,19 +95,22 @@ if uploaded_file is not None:
                             truncation=True
                         ).to(model.device)
 
-                        # Añadir atención y ejecutar la generación en lotes
-                        with torch.cuda.amp.autocast():  # Mixed Precision para más velocidad
+                        with torch.cuda.amp.autocast():  # Usar Mixed Precision
                             model_inputs['attention_mask'] = (model_inputs['input_ids'] != tokenizer.pad_token_id).int()
                             generated_ids = model.generate(
                                 **model_inputs, 
-                                max_new_tokens=50, 
+                                max_new_tokens=20,  # Reducir para minimizar el uso de memoria
                                 num_beams=1  # Desactivar búsqueda en beam para más velocidad
                             )
-                        
+
                         # Decodificar el resultado y añadirlo a la lista de resultados
                         decoded = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
                         all_scores.extend([0.95] * len(batch))  # Simulación para demostración
-                    
+
+                        # Liberar memoria después de cada batch
+                        del model_inputs, generated_ids
+                        torch.cuda.empty_cache()
+
                     # Asignar puntajes al DataFrame
                     df['Score'] = all_scores
 
@@ -142,7 +123,6 @@ if uploaded_file is not None:
         st.error("La columna 'job_title' no se encuentra en el archivo CSV.")
 
 
-
 '''