Update app.py

app.py

@@ -14,23 +14,21 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
 tokenizer.pad_token = tokenizer.eos_token
 
-MAX_INPUT_TOKEN_LENGTH = 10000
+MAX_INPUT_TOKEN_LENGTH = 4096
 
-def generate_response(input_text, temperature=0.5, max_new_tokens=100):
+def generate_response(input_text, temperature=0.5, max_new_tokens=20):
     input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
 
     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.")
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         input_ids=input_ids,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_k=50,
-        top_p=0.9,
+        num_beams=3,  # Usar beam search
         temperature=temperature,
         eos_token_id=[tokenizer.eos_token_id]
     )
@@ -38,14 +36,17 @@ def generate_response(input_text, temperature=0.5, max_new_tokens=100):
     try:
         t = Thread(target=model.generate, kwargs=generate_kwargs)
         t.start()
-        t.join()  # Esperar a que el hilo termine
+        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.")
-        return "".join(outputs)
+        
+        # Post-procesamiento más restrictivo
+        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."
@@ -65,13 +66,13 @@ def main():
             st.write("Archivo CSV cargado exitosamente:")
             st.write(df.head())
 
-            initial_prompt = f"I have a list of job titles: {job_titles}. Please extract and return only the first job title from this list without repeating."
+            initial_prompt = f"The list of job titles is: {job_titles}. Extract only the first job title from the list and return it as the answer."
             st.write(f"Query: {query}")
             st.write(f"Prompt inicial: {initial_prompt}")
 
             if st.button("Generar respuesta"):
                 with st.spinner("Generando respuesta..."):
-                    response = generate_response(initial_prompt, temperature=0.2)
+                    response = generate_response(initial_prompt, temperature=0.5)
                     if response:
                         st.write(f"Respuesta del modelo: {response}")
                     else: