Update app.py

app.py

@@ -15,35 +15,42 @@ tokenizer, yi_coder_model, yi_coder_device = load_yi_coder_model()
 # Conectar a Pinecone
 index = connect_to_pinecone()
 
-# Función para generar código con Yi-Coder
+# Función para generar código utilizando Yi-Coder
 @gpu_decorator(duration=100)
 def generate_code(system_prompt, user_prompt, max_length):
     device = yi_coder_device
     model = yi_coder_model
-    tokenizer_ = tokenizer
-
-    # Combina el system_prompt y user_prompt sin formato de chat
-    prompt = f"{system_prompt}\n{user_prompt}"
-
-    # Tokeniza el prompt
-    model_inputs = tokenizer_(prompt, return_tensors="pt").to(device)
+    tokenizer_ = tokenizer  # Ya lo tenemos cargado
+
+    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
+    )
+    model_inputs = tokenizer_([text], return_tensors="pt").to(device)
 
     with torch.no_grad():
-        # Genera la respuesta
         generated_ids = model.generate(
             model_inputs.input_ids,
             max_new_tokens=max_length,
-            eos_token_id=tokenizer_.eos_token_id
+            eos_token_id=tokenizer_.eos_token_id  
         )
 
-    # Decodifica el texto generado
-    generated_text = tokenizer_.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-    # Extrae solo la parte generada después del prompt inicial
-    response = generated_text[len(prompt):].strip()
+    # 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)
+    ]
 
+    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):