File size: 3,041 Bytes
751e7d4
085880a
f7a9983
 
085880a
4b143e6
f7a9983
 
 
751e7d4
 
71c83be
77de53d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
751e7d4
a2ca183
f7a9983
 
 
 
 
 
4a6d7a3
 
f7a9983
 
a2ca183
f7a9983
 
085880a
6772cf6
 
a29437c
 
 
 
 
 
6772cf6
 
 
a29437c
 
 
 
6772cf6
 
 
085880a
6772cf6
1cb1025
085880a
 
170063c
085880a
f7a9983
 
1cb1025
 
 
 
 
 
 
 
 
 
 
 
 
 
d4b2751
1cb1025
085880a
 
6159031
1cb1025
085880a
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import os
import gradio as gr
from huggingface_hub import InferenceClient
from e2b_code_interpreter import Sandbox

from utils import run_interactive_notebook

message_history = None

E2B_API_KEY = os.environ['E2B_API_KEY']
HF_TOKEN = os.environ['HF_TOKEN']
DEFAULT_MAX_TOKENS = 512
DEFAULT_SYSTEM_PROMPT = """Environment: ipython

You are a code assistant with access to a ipython interpreter.
You solve tasks step-by-step and rely on code execution results.
Don't make any assumptions about the data but always check the format first.
If you generate code in your response you always run it in the interpreter.
When fix a mistake in the code always run it again.

Follow these steps given a new task and dataset:
1. Read in the data and make sure you understand each files format and content by printing useful information. 
2. Execute the code at this point and don't try to write a solution before looking at the execution result.
3. After exploring the format write a quick action plan to solve the task from the user.
4. Then call the ipython interpreter directly with the solution and look at the execution result.
5. If there is an issue with the code, reason about potential issues and then propose a solution and execute again the fixed code and check the result.
Always run the code at each step and repeat the steps if necessary until you reach a solution. 

NEVER ASSUME, ALWAYS VERIFY!"""


def execute_jupyter_agent(sytem_prompt, user_input, max_new_tokens):
    client = InferenceClient(api_key=HF_TOKEN)
    model = "meta-llama/Llama-3.1-8B-Instruct"

    sbx = Sandbox(api_key=E2B_API_KEY)

    messages = [
        {"role": "system", "content": sytem_prompt},
        {"role": "user", "content": user_input}
    ]

    for notebook_html, messages in run_interactive_notebook(client, model, messages, sbx, max_new_tokens=max_new_tokens):
        message_history = messages
        yield notebook_html


css = """
#component-0 {
    height: 100vh;
    overflow-y: auto;
    padding: 20px;
}

.gradio-container {
    height: 100vh !important;
}

.contain {
    height: 100vh !important;
}
"""


# Create the interface
with gr.Blocks(css=css) as demo:
    gr.Markdown("# Jupyter Agent!")
    
    with gr.Row():
        user_input = gr.Textbox(label="User prompt", value="Solve the Lotka-Volterra equation and plot the results.", lines=3)
    
    generate_btn = gr.Button("Let's go!")
    output = gr.HTML(label="Jupyter Notebook")

    with gr.Accordion("Advanced Settings", open=False):
        system_input = gr.Textbox(
            label="System Prompt",
            value=DEFAULT_SYSTEM_PROMPT,
            elem_classes="input-box"
        )

        max_tokens = gr.Number(
            label="Max New Tokens",
            value=DEFAULT_MAX_TOKENS,
            minimum=128,
            maximum=2048,
            step=8,
            interactive=True
        )
    
    generate_btn.click(
        fn=execute_jupyter_agent,
        inputs=[system_input, user_input, max_tokens],
        outputs=output
    )

demo.launch()