Commit

git commit .

streaming.py

@@ -1,6 +1,6 @@
 from transformers.agents.agent_types import AgentAudio, AgentImage, AgentText, AgentType
 from transformers.agents import ReactAgent
-
+import spaces
 
 def pull_message(step_log: dict):
     try:
@@ -29,7 +29,7 @@ def pull_message(step_log: dict):
             metadata={"title": "💥 Error"},
         )
 
-
+@spaces.GPU
 def stream_to_gradio(agent: ReactAgent, task: str, **kwargs):
     """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
 

test_app.py

@@ -0,0 +1,133 @@
+import os
+import shutil
+import gradio as gr
+from transformers import ReactCodeAgent, HfEngine, Tool
+import pandas as pd
+
+from gradio import Chatbot
+from test_streaming import stream_to_gradio
+from huggingface_hub import login
+from gradio.data_classes import FileData
+
+login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
+
+llm_engine = HfEngine("meta-llama/Meta-Llama-3.1-70B-Instruct")
+
+agent = ReactCodeAgent(
+    tools=[],
+    llm_engine=llm_engine,
+    additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn", "scipy.stats"],
+    max_iterations=10,
+)
+
+base_prompt = """You are an expert data analyst.
+According to the features you have and the data structure given below, determine which feature should be the target.
+Then list 3 interesting questions that could be asked on this data, for instance about specific correlations with target variable.
+Then answer these questions one by one, by finding the relevant numbers.
+Meanwhile, plot some figures using matplotlib/seaborn and save them to the (already existing) folder './figures/': take care to clear each figure with plt.clf() before doing another plot.
+
+In your final answer: summarize these correlations and trends
+After each number derive real worlds insights, for instance: "Correlation between is_december and boredness is 1.3453, which suggest people are more bored in winter".
+Your final answer should be a long string with at least 3 numbered and detailed parts.
+
+Structure of the data:
+{structure_notes}
+
+The data file is passed to you as the variable data_file, it is a pandas dataframe, you can use it directly.
+DO NOT try to load data_file, it is already a dataframe pre-loaded in your python interpreter!
+"""
+
+example_notes="""This data is about the Titanic wreck in 1912.
+The target figure is the survival of passengers, notes by 'Survived'
+pclass: A proxy for socio-economic status (SES)
+1st = Upper
+2nd = Middle
+3rd = Lower
+age: Age is fractional if less than 1. If the age is estimated, is it in the form of xx.5
+sibsp: The dataset defines family relations in this way...
+Sibling = brother, sister, stepbrother, stepsister
+Spouse = husband, wife (mistresses and fiancés were ignored)
+parch: The dataset defines family relations in this way...
+Parent = mother, father
+Child = daughter, son, stepdaughter, stepson
+Some children travelled only with a nanny, therefore parch=0 for them."""
+
+def get_images_in_directory(directory):
+    image_extensions = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.tiff'}
+
+    image_files = []
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if os.path.splitext(file)[1].lower() in image_extensions:
+                image_files.append(os.path.join(root, file))
+    return image_files
+
+def interact_with_agent(file_input, additional_notes):
+    shutil.rmtree("./figures")
+    os.makedirs("./figures")
+
+    data_file = pd.read_csv(file_input)
+    data_structure_notes = f"""- Description (output of .describe()):
+    {data_file.describe()}
+    - Columns with dtypes:
+    {data_file.dtypes}"""
+
+    prompt = base_prompt.format(structure_notes=data_structure_notes)
+
+    if additional_notes and len(additional_notes) > 0:
+        prompt += "\nAdditional notes on the data:\n" + additional_notes
+
+    messages = [gr.ChatMessage(role="user", content=prompt)]
+    yield messages + [
+        gr.ChatMessage(role="assistant", content="⏳ _Starting task..._")
+    ]
+
+    plot_image_paths = {}
+    for msg in stream_to_gradio(agent, prompt, data_file=data_file):
+        messages.append(msg)
+        for image_path in get_images_in_directory("./figures"):
+            if image_path not in plot_image_paths:
+                image_message = gr.ChatMessage(
+                    role="assistant",
+                    content=FileData(path=image_path, mime_type="image/png"),
+                )
+                plot_image_paths[image_path] = True
+                messages.append(image_message)
+        yield messages + [
+            gr.ChatMessage(role="assistant", content="⏳ _Still processing..._")
+        ]
+    yield messages
+
+
+with gr.Blocks(
+    theme=gr.themes.Soft(
+        primary_hue=gr.themes.colors.yellow,
+        secondary_hue=gr.themes.colors.blue,
+    )
+) as demo:
+    gr.Markdown("""# Llama-3.1 Data analyst 📊🤔
+
+Drop 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!**""")
+    file_input = gr.File(label="Your file to analyze")
+    text_input = gr.Textbox(
+        label="Additional notes to support the analysis"
+    )
+    submit = gr.Button("Run analysis!", variant="primary")
+    chatbot = gr.Chatbot(
+        label="Data Analyst Agent",
+        type="messages",
+        avatar_images=(
+            None,
+            "https://em-content.zobj.net/source/twitter/53/robot-face_1f916.png",
+        ),
+    )
+    gr.Examples(
+        examples=[["./example/titanic.csv", example_notes]],
+        inputs=[file_input, text_input],
+        cache_examples=False
+    )
+
+    submit.click(interact_with_agent, [file_input, text_input], [chatbot])
+
+if __name__ == "__main__":
+    demo.launch()

test_streaming.py

@@ -0,0 +1,64 @@
+from transformers.agents.agent_types import AgentAudio, AgentImage, AgentText, AgentType
+from transformers.agents import ReactAgent
+
+
+def pull_message(step_log: dict):
+    try:
+        from gradio import ChatMessage
+    except ImportError:
+        raise ImportError("Gradio should be installed in order to launch a gradio demo.")
+
+    if step_log.get("rationale"):
+        yield ChatMessage(role="assistant", content=step_log["rationale"])
+    if step_log.get("tool_call"):
+        used_code = step_log["tool_call"]["tool_name"] == "code interpreter"
+        content = step_log["tool_call"]["tool_arguments"]
+        if used_code:
+            content = f"```py\n{content}\n```"
+        yield ChatMessage(
+            role="assistant",
+            metadata={"title": f"🛠️ Used tool {step_log['tool_call']['tool_name']}"},
+            content=content,
+        )
+    if step_log.get("observation"):
+        yield ChatMessage(role="assistant", content=f"```\n{step_log['observation']}\n```")
+    if step_log.get("error"):
+        yield ChatMessage(
+            role="assistant",
+            content=str(step_log["error"]),
+            metadata={"title": "💥 Error"},
+        )
+
+
+def stream_to_gradio(agent: ReactAgent, task: str, **kwargs):
+    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""
+
+    try:
+        from gradio import ChatMessage
+    except ImportError:
+        raise ImportError("Gradio should be installed in order to launch a gradio demo.")
+
+    class Output:
+        output: AgentType | str = None
+
+    for step_log in agent.run(task, stream=True, **kwargs):
+        if isinstance(step_log, dict):
+            for message in pull_message(step_log):
+                print("message", message)
+                yield message
+
+    Output.output = step_log
+    if isinstance(Output.output, AgentText):
+        yield ChatMessage(role="assistant", content=f"**Final answer:**\n```\n{Output.output.to_string()}\n```")
+    elif isinstance(Output.output, AgentImage):
+        yield ChatMessage(
+            role="assistant",
+            content={"path": Output.output.to_string(), "mime_type": "image/png"},
+        )
+    elif isinstance(Output.output, AgentAudio):
+        yield ChatMessage(
+            role="assistant",
+            content={"path": Output.output.to_string(), "mime_type": "audio/wav"},
+        )
+    else:
+        yield ChatMessage(role="assistant", content=Output.output)