Update app.py

app.py

@@ -15,31 +15,31 @@ tokenizer, yi_coder_model, yi_coder_device = load_yi_coder_model()
 # Conectar a Pinecone
 index = connect_to_pinecone()
 
-# Función para generar código utilizando Yi-Coder
+# Función para generar código con 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  # Ya lo tenemos cargado
+    tokenizer_ = tokenizer
 
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_prompt}
-    ]
+    # Combina el system_prompt y user_prompt sin formato de chat
+    prompt = f"{system_prompt}\n{user_prompt}"
 
-    # Preparamos el input para el modelo
-    prompt = system_prompt + "\n" + user_prompt
+    # Tokeniza el prompt
     model_inputs = tokenizer_(prompt, 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
         )
 
-    # Extraer solo la parte generada
+    # 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()
 
     return response