Bl4ck0z/IA-ChatBot-CSV

1import os2import shutil3import gradio as gr4from transformers import ReactCodeAgent, HfApiEngine, Tool5import pandas as pd6 7from gradio import Chatbot8from transformers.agents import stream_to_gradio9from huggingface_hub import login10from gradio.data_classes import FileData11 12# Iniciar sesión en Hugging Face con el token desde las variables de entorno13login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))14 15# Configuración del motor LLM16llm_engine = HfApiEngine("meta-llama/Llama-3.3-70B-Instruct")17 18# Configuración del agente de análisis19agent = ReactCodeAgent(20    tools=[],21    llm_engine=llm_engine,22    additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn", "scipy.stats"],23    max_iterations=10,24)25 26# Prompt base para el análisis27base_prompt = """Eres un analista de datos experto.28Según las características disponibles y la estructura de datos descrita a continuación, determina cuál debe ser la variable objetivo.29Luego, formula 3 preguntas interesantes que podrían hacerse sobre estos datos, por ejemplo, sobre correlaciones específicas con la variable objetivo.30Después responde estas preguntas una por una, encontrando los números relevantes.31Mientras tanto, genera algunas figuras utilizando matplotlib/seaborn y guárdalas en la carpeta existente './figures/': recuerda limpiar cada figura con plt.clf() antes de generar otra.32 33En tu respuesta final: resume estas correlaciones y tendencias.34Después de cada número, proporciona conclusiones del mundo real, por ejemplo: "La correlación entre is_december y boredness es 1.3453, lo que sugiere que las personas están más aburridas en invierno".35Tu respuesta final debe ser un texto largo con al menos 3 partes numeradas y detalladas.36 37Estructura de los datos:38{structure_notes}39 40El archivo de datos se pasa como la variable data_file, ya es un dataframe de pandas, ¡puedes usarlo directamente!41NO intentes cargar data_file, ya está pre-cargado en tu intérprete de Python.42"""43 44# Notas de ejemplo45example_notes = """Estos datos son sobre el naufragio del Titanic en 1912.46La variable objetivo es la supervivencia de los pasajeros, indicada por 'Survived'.47pclass: Un proxy para el estatus socioeconómico (SES)481st = Alta492nd = Media503rd = Baja51age: Edad. Es fraccionaria si es menor a 1. Si la edad es estimada, está en forma xx.5.52sibsp: El dataset define las relaciones familiares así...53Hermano = hermano, hermana, hermanastro, hermanastra54Cónyuge = esposo, esposa (amantes y prometidos fueron ignorados)55parch: El dataset define las relaciones familiares así...56Padre = madre, padre57Hijo = hija, hijo, hijastra, hijastro58Algunos niños viajaron solo con una niñera, por lo tanto parch=0 para ellos."""59 60def get_images_in_directory(directory):61    image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff'}62 63    image_files = []64    for root, dirs, files in os.walk(directory):65        for file in files:66            if os.path.splitext(file)[1].lower() in image_extensions:67                image_files.append(os.path.join(root, file))68    return image_files69 70def interact_with_agent(file_input, additional_notes):71    shutil.rmtree("./figures")72    os.makedirs("./figures")73 74    data_file = pd.read_csv(file_input)75    data_structure_notes = f"""- Description (output of .describe()):76    {data_file.describe()}77    - Columns with dtypes:78    {data_file.dtypes}"""79 80    prompt = base_prompt.format(structure_notes=data_structure_notes)81 82    if additional_notes and len(additional_notes) > 0:83        prompt += "\nAdditional notes on the data:\n" + additional_notes84 85    messages = [gr.ChatMessage(role="user", content=prompt)]86    yield messages + [87        gr.ChatMessage(role="assistant", content="⏳ _Starting task..._")88    ]89 90    plot_image_paths = {}91    for msg in stream_to_gradio(agent, prompt, data_file=data_file):92        messages.append(msg)93        for image_path in get_images_in_directory("./figures"):94            if image_path not in plot_image_paths:95                image_message = gr.ChatMessage(96                    role="assistant",97                    content=FileData(path=image_path, mime_type="image/png"),98                )99                plot_image_paths[image_path] = True100                messages.append(image_message)101        yield messages + [102            gr.ChatMessage(role="assistant", content="⏳ _Still processing..._")103        ]104    yield messages105 106 107with gr.Blocks(108    theme=gr.themes.Soft(109        primary_hue=gr.themes.colors.yellow,110        secondary_hue=gr.themes.colors.blue,111    )112) as demo:113    gr.Markdown("""# Llama-3.1 Data analyst 📊🤔114Drop a `.csv` file below, add notes to describe this data if needed, and **Llama-3.1-70B will analyze the file content and draw figures for you!**""")115    file_input = gr.File(label="Your file to analyze")116    text_input = gr.Textbox(117        label="Additional notes to support the analysis"118    )119    submit = gr.Button("Run analysis!", variant="primary")120    chatbot = gr.Chatbot(121        label="Data Analyst Agent",122        type="messages",123        avatar_images=(124            None,125            "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png",126        ),127    )128    gr.Examples(129        examples=[["./example/titanic.csv", example_notes]],130        inputs=[file_input, text_input],131        cache_examples=False132    )133 134    submit.click(interact_with_agent, [file_input, text_input], [chatbot])135 136if __name__ == "__main__":137    demo.launch()