Update app.py

app.py

@@ -16,7 +16,7 @@ tokenizer.pad_token = tokenizer.eos_token
 
 MAX_INPUT_TOKEN_LENGTH = 4096
 
-def generate_response(input_text, temperature=0.5, max_new_tokens=20):
+def generate_response(input_text, temperature=0.5, max_new_tokens=50):
     input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
 
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
@@ -28,7 +28,9 @@ def generate_response(input_text, temperature=0.5, max_new_tokens=20):
         input_ids=input_ids,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        num_beams=3,  # Usar beam search
+        do_sample=True,
+        top_k=40,
+        top_p=0.9,
         temperature=temperature,
         eos_token_id=[tokenizer.eos_token_id]
     )
@@ -44,7 +46,6 @@ def generate_response(input_text, temperature=0.5, max_new_tokens=20):
         if not outputs:
             raise ValueError("No se generó ninguna respuesta.")
         
-        # Post-procesamiento más restrictivo
         response = "".join(outputs).strip().split("\n")[0]
         return response
     except Exception as e:
@@ -61,14 +62,25 @@ def main():
         
         if 'job_title' in df.columns:
             job_titles = df['job_title'].tolist()
-            query = "aspiring human resources specialist"
-            
-            st.write("Archivo CSV cargado exitosamente:")
-            st.write(df.head())
 
-            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}")
+            # Definir el prompt con in-context learning
+            initial_prompt = (
+                "Here are some examples of job title extraction:\n"
+                "Example 1:\n"
+                "List: ['Data Scientist', 'Machine Learning Engineer', 'AI Researcher']\n"
+                "First job title: 'Data Scientist'\n"
+                "\n"
+                "Example 2:\n"
+                "List: ['Software Developer', 'Backend Engineer', 'Frontend Developer']\n"
+                "First job title: 'Software Developer'\n"
+                "\n"
+                "Now, extract the first job title from the following list:\n"
+                f"List: {job_titles}\n"
+                "First job title:"
+            )
+            
+            st.write("Prompt inicial con In-context Learning:")
+            st.write(initial_prompt)
 
             if st.button("Generar respuesta"):
                 with st.spinner("Generando respuesta..."):
@@ -89,3 +101,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+