Update app.py

app.py

@@ -16,42 +16,38 @@ tokenizer.pad_token = tokenizer.eos_token
 
 MAX_INPUT_TOKEN_LENGTH = 10000
 
-def generate_response(input_text, temperature=0.2, max_new_tokens=20):
-    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
+# Asegurar que el token de padding esté configurado
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token  # Asignar el token de padding al token de fin de oración
+
+def generate_response(input_text, max_new_tokens=50, temperature=0.5):
+    """
+    Genera una respuesta usando el modelo de lenguaje con máscara de atención.
+    """
+    # Tokenizar la entrada y crear la máscara de atención
+    inputs = tokenizer(
+        input_text,
+        return_tensors='pt',
+        padding=True,
+        truncation=True,
+        max_length=512  # Ajustar según sea necesario
+    )
 
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        st.warning(f"Se recortó la entrada porque excedió el límite de {MAX_INPUT_TOKEN_LENGTH} tokens.")
+    input_ids = inputs['input_ids'].to(model.device)
+    attention_mask = inputs['attention_mask'].to(model.device)
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=120.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        streamer=streamer,
+    # Generar texto con la máscara de atención y el token de padding
+    outputs = model.generate(
+        input_ids,
+        attention_mask=attention_mask,
         max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_k=20,
-        top_p=0.9,
         temperature=temperature,
-        num_beams= 3,
-        eos_token_id=[tokenizer.eos_token_id]
+        pad_token_id=tokenizer.pad_token_id  # Usar el token de padding configurado
     )
 
-    try:
-        t = Thread(target=model.generate, kwargs=generate_kwargs)
-        t.start()
-        t.join()  # Asegura que la generación haya terminado
-
-        outputs = []
-        for text in streamer:
-            outputs.append(text)
-        if not outputs:
-            raise ValueError("No se generó ninguna respuesta.")
-        
-        response = "".join(outputs).strip().split("\n")[0]
-        return response
-    except Exception as e:
-        st.error(f"Error durante la generación: {e}")
-        return "Error en la generación de texto."
+    # Decodificar la respuesta generada
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 def main():
     st.title("Chat con Meta Llama 3.2 1B")