Update app.py

app.py

@@ -1,9 +1,9 @@
 import pandas as pd
 import streamlit as st
-from langchain_huggingface import HuggingFacePipeline  # Nueva importación
+from langchain_huggingface import HuggingFacePipeline
 from langchain_core.prompts import PromptTemplate
 from langchain.chains import LLMChain
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, pipeline
 from huggingface_hub import login
 import torch
 import json
@@ -13,15 +13,25 @@ from datetime import datetime
 huggingface_token = st.secrets["FIREWORKS"]
 login(huggingface_token)
 
-# Configurar modelo Fireworks desde Hugging Face
+# Configurar modelo Fireworks con cuantización int8
+quant_config = BitsAndBytesConfig.from_model_type(
+    "int8",  # Cuantización para reducir el tamaño y acelerar
+    quantization_scheme="gptq"
+)
+
 model_id = "fireworks-ai/firefunction-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id, 
     device_map="auto", 
-    torch_dtype=torch.float16
+    torch_dtype=torch.float16,
+    quantization_config=quant_config
 )
 
+# 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 = [
     {
@@ -50,7 +60,7 @@ fireworks_pipeline = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_new_tokens=128
+    max_new_tokens=50  # Reducir max_new_tokens para acelerar
 )
 
 # Adaptar el pipeline a LangChain
@@ -69,6 +79,10 @@ if uploaded_file is not None:
         query = 'aspiring human resources specialist'
         job_titles = df['job_title'].tolist()
 
+        # Procesar en lotes para optimización
+        batch_size = 16  # Ajusta según la memoria de la GPU
+        job_titles_batches = [job_titles[i:i+batch_size] for i in range(0, len(job_titles), batch_size)]
+
         # Definir el prompt para Fireworks
         prompt_template = PromptTemplate(
             template=(
@@ -88,29 +102,32 @@ if uploaded_file is not None:
         # Ejecutar la generación con Fireworks y funciones
         if st.button("Calcular Similitud de Coseno"):
             with st.spinner("Calculando similitudes con Fireworks..."):
+                all_scores = []
                 try:
-                    # Preparar mensajes y funciones para Fireworks
-                    messages = [
-                        {'role': 'system', 'content': 'You are a helpful assistant with access to functions. Use them if required.'},
-                        {'role': 'user', 'content': f'Calculate cosine similarity for query: {query} with job titles.'}
-                    ]
-                    now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-
-                    model_inputs = tokenizer.apply_chat_template(
-                        messages,
-                        functions=functions,
-                        datetime=now,
-                        return_tensors="pt"
-                    ).to(model.device)
-
-                    # Generar resultados con Fireworks
-                    generated_ids = model.generate(model_inputs, max_new_tokens=128)
-                    decoded = tokenizer.batch_decode(generated_ids)
-                    st.write("Respuesta del modelo:")
-                    st.write(decoded[0])
-
-                    # Simular la asignación de puntajes en la columna 'Score' (basado en la respuesta del modelo)
-                    df['Score'] = [0.95] * len(df)  # Simulación para la demostración
+                    for batch in job_titles_batches:
+                        # Tokenizar la entrada con atención en lotes
+                        model_inputs = tokenizer(
+                            batch,
+                            return_tensors="pt",
+                            padding=True,
+                            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
+                            model_inputs['attention_mask'] = (model_inputs['input_ids'] != tokenizer.pad_token_id).int()
+                            generated_ids = model.generate(
+                                **model_inputs, 
+                                max_new_tokens=50, 
+                                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
+                    
+                    # Asignar puntajes al DataFrame
+                    df['Score'] = all_scores
 
                     # Mostrar el dataframe actualizado
                     st.write("DataFrame con los puntajes de similitud:")
@@ -121,6 +138,7 @@ if uploaded_file is not None:
         st.error("La columna 'job_title' no se encuentra en el archivo CSV.")
 
 
+
 '''