Update app.py

app.py

@@ -3,20 +3,18 @@
 import gradio as gr
 import os
 from models import load_embedding_model, load_yi_coder_model
-from pinecone_utils import connect_to_pinecone, vector_search  # Ahora debería funcionar correctamente
+from pinecone_utils import connect_to_pinecone, vector_search
 from ui import build_interface
 from config import SIMILARITY_THRESHOLD_DEFAULT, SYSTEM_PROMPT, MAX_LENGTH_DEFAULT
 from decorators import gpu_decorator
 import torch
 
 ########################
-
 from utils import process_tags_chat
-
 ########################
 
-
 # Cargar modelos
+# Estas funciones ahora devuelven los modelos y dispositivos correctamente configurados.
 embedding_model = load_embedding_model()
 tokenizer, yi_coder_model, yi_coder_device = load_yi_coder_model()
 
@@ -26,21 +24,22 @@ index = connect_to_pinecone()
 # Función para generar código utilizando Yi-Coder
 @gpu_decorator(duration=100)
 def generate_code(system_prompt, user_prompt, max_length):
+    # CORRECTO: yi_coder_device ya es un objeto torch.device, no una cadena.
     device = yi_coder_device
     model = yi_coder_model
-    tokenizer_ = tokenizer  # Ya lo tenemos cargado
+    tokenizer_ = tokenizer
 
     messages = [
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_prompt}
     ]
     
-    # Aplicar la plantilla de chat y preparar el texto
     text = tokenizer_.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
     )
+    # .to(device) funcionará sin problemas porque device es un objeto torch.device.
     model_inputs = tokenizer_([text], return_tensors="pt").to(device)
 
     with torch.no_grad():
@@ -50,7 +49,6 @@ def generate_code(system_prompt, user_prompt, max_length):
             eos_token_id=tokenizer_.eos_token_id  
         )
 
-    # Extraer solo la parte generada
     generated_ids = [
         output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
     ]
@@ -58,19 +56,16 @@ def generate_code(system_prompt, user_prompt, max_length):
     response = tokenizer_.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return response
 
-
 # Función para combinar búsqueda vectorial y Yi-Coder
 @gpu_decorator(duration=100)
 def combined_function(user_prompt, similarity_threshold, selected_option, system_prompt, max_length):
     def get_partial_message(response):
-        """Obtiene el contenido después de 'Respuesta:' si está presente en la respuesta."""
         if "Respuesta:" in response:
-            return response.split("Respuesta:")[1].strip()  # Tomar solo el texto después de 'Respuesta:'
+            return response.split("Respuesta:")[1].strip()
         else:
-            return response  # Devolver la respuesta completa si no contiene 'Respuesta:'
+            return response
 
     if selected_option == "Solo Búsqueda Vectorial":
-        # Realizar búsqueda vectorial
         search_results = vector_search(user_prompt, embedding_model, index)
         if search_results:
             content = search_results[0]['content']
@@ -79,12 +74,10 @@ def combined_function(user_prompt, similarity_threshold, selected_option, system
         else:
             return "No se encontraron resultados en Pinecone.", None
     elif selected_option == "Solo Yi-Coder":
-        # Generar respuesta usando Yi-Coder
         yi_coder_response = generate_code(system_prompt, user_prompt, max_length)
         partial_message = get_partial_message(yi_coder_response)
         return partial_message, None
     elif selected_option == "Ambos (basado en umbral de similitud)":
-        # Realizar búsqueda vectorial
         search_results = vector_search(user_prompt, embedding_model, index)
         if search_results:
             top_result = search_results[0]
@@ -103,7 +96,6 @@ def combined_function(user_prompt, similarity_threshold, selected_option, system
     else:
         return "Opción no válida.", None
 
-
 # Funciones para el procesamiento de entradas y actualización de imágenes
 def process_input(message, history, selected_option, similarity_threshold, system_prompt, max_length):
     response, image = combined_function(message, similarity_threshold, selected_option, system_prompt, max_length)
@@ -111,15 +103,6 @@ def process_input(message, history, selected_option, similarity_threshold, syste
     return history, history, image
 
 def update_image(image_url):
-    """
-    Retorna los datos binarios de la imagen para ser mostrados en Gradio.
-
-    Args:
-        image_url (str): Ruta de la imagen.
-
-    Returns:
-        bytes o None: Datos binarios de la imagen si existe, de lo contrario None.
-    """
     if image_url and os.path.exists(image_url):
         try:
             with open(image_url, "rb") as img_file:
@@ -132,7 +115,6 @@ def update_image(image_url):
         print("No se encontró una imagen válida.")
         return None
 
-
 def send_preset_question(question, history, selected_option, similarity_threshold, system_prompt, max_length):
     return process_input(question, history, selected_option, similarity_threshold, system_prompt, max_length)
 
@@ -140,4 +122,4 @@ def send_preset_question(question, history, selected_option, similarity_threshol
 demo = build_interface(process_input, send_preset_question, update_image)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()