update

app.py

@@ -1,11 +1,12 @@
 import os
 import gradio as gr
 from gradio.utils import get_space
-from huggingface_hub import InferenceClient
 from e2b_code_interpreter import Sandbox
 from pathlib import Path
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import json
+import re
+import torch
 
 if not get_space():
     try:
@@ -20,6 +21,8 @@ from utils import (
     run_interactive_notebook,
     create_base_notebook,
     update_notebook_display,
+    update_notebook_with_cell,
+    update_notebook_with_markdown,
 )
 
 E2B_API_KEY = os.environ["E2B_API_KEY"]
@@ -37,9 +40,47 @@ with open(TMP_DIR+"jupyter-agent.ipynb", 'w', encoding='utf-8') as f:
 with open("ds-system-prompt.txt", "r") as f:
     DEFAULT_SYSTEM_PROMPT = f.read()
 
+# Add this constant at the top with other constants
+MAX_TURNS = 10
+
+# Replace the client initialization with local model loading
+def load_model_and_tokenizer(model_name="bigcomputer/jupycoder-7b-lora-350"):
+    if model_name == "bigcomputer/jupycoder-7b-lora-350":
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return model, tokenizer
+
+# Function to extract code and text from model response
+def parse_model_response(response_text):
+    cells = []
+    # Split by code blocks
+    parts = re.split(r'(```python[\s\S]*?```)', response_text)
+    
+    for part in parts:
+        if part.strip():
+            if part.startswith('```python'):
+                # Extract code without the markers
+                code = re.sub(r'```python\n|```', '', part).strip()
+                cells.append({"type": "code", "content": code})
+            else:
+                # Regular text becomes markdown
+                cells.append({"type": "markdown", "content": part.strip()})
+    
+    return cells
 
 def execute_jupyter_agent(
-    sytem_prompt, user_input, max_new_tokens, model, files, message_history, request: gr.Request
+    system_prompt, user_input, max_new_tokens, model_name, files, message_history, request: gr.Request
 ):
     if request.session_hash not in SANDBOXES:
         SANDBOXES[request.session_hash] = Sandbox(api_key=E2B_API_KEY)
@@ -49,11 +90,9 @@ def execute_jupyter_agent(
     os.makedirs(save_dir, exist_ok=True)
     save_dir = os.path.join(save_dir, 'jupyter-agent.ipynb')
 
-    client = InferenceClient(api_key=HF_TOKEN)
-
-    tokenizer = AutoTokenizer.from_pretrained(model)
-    # model = "meta-llama/Llama-3.1-8B-Instruct"
+    model, tokenizer = load_model_and_tokenizer(model_name)
 
+    # Handle file uploads
     filenames = []
     if files is not None:
         for filepath in files:
@@ -63,28 +102,73 @@ def execute_jupyter_agent(
                 sbx.files.write(filpath.name, file)
                 filenames.append(filpath.name)
 
-    # Initialize message_history if it doesn't exist
+    # Initialize conversation
     if len(message_history) == 0:
-        message_history.append(
-            {
-                "role": "system",
-                "content": sytem_prompt.format("- " + "\n- ".join(filenames)),
-            }
-        )
+        message_history.append({
+            "role": "system",
+            "content": system_prompt.format("- " + "\n- ".join(filenames))
+        })
     message_history.append({"role": "user", "content": user_input})
 
-    print("history:", message_history)
+    # Create initial notebook
+    notebook_data = create_base_notebook([])
+    turn_count = 0
 
-    for notebook_html, notebook_data, messages in run_interactive_notebook(
-        client, model, tokenizer, message_history, sbx, max_new_tokens=max_new_tokens
-    ):
-        message_history = messages
+    while turn_count < MAX_TURNS:
+        turn_count += 1
         
-        yield notebook_html, message_history, TMP_DIR+"jupyter-agent.ipynb"
-    
+        # Generate response
+        input_text = "\n".join([msg["content"] for msg in message_history])
+        inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+        
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+        )
+        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # Parse response into cells
+        cells = parse_model_response(response_text)
+        
+        # Process each cell
+        has_code = False
+        for cell in cells:
+            if cell["type"] == "code":
+                has_code = True
+                # Execute code cell
+                result = sbx.python.run(cell["content"])
+                # Add code cell and output to notebook
+                notebook_data = update_notebook_with_cell(notebook_data, cell["content"], result)
+                # Add execution result to message history
+                message_history.append({
+                    "role": "assistant",
+                    "content": cell["content"]
+                })
+                message_history.append({
+                    "role": "user",
+                    "content": f"Execution result:\n{result}"
+                })
+            else:
+                # Add markdown cell to notebook
+                notebook_data = update_notebook_with_markdown(notebook_data, cell["content"])
+                message_history.append({
+                    "role": "assistant",
+                    "content": cell["content"]
+                })
+            
+            # Update display after each cell
+            notebook_html = update_notebook_display(notebook_data)
+            yield notebook_html, message_history, save_dir
+
+        # If no code was generated or we've reached max turns, stop
+        if not has_code or turn_count >= MAX_TURNS:
+            break
+
+    # Save final notebook
     with open(save_dir, 'w', encoding='utf-8') as f:
         json.dump(notebook_data, f, indent=2)
-    yield notebook_html, message_history, save_dir
 
 def clear(msg_state):
     msg_state = []
@@ -145,11 +229,10 @@ with gr.Blocks() as demo:
             )
 
             model = gr.Dropdown(
-                value="meta-llama/Llama-3.1-8B-Instruct",
+                value="bigcomputer/jupycoder-7b-lora-350",
                 choices=[
-                    "meta-llama/Llama-3.2-3B-Instruct",
-                    "meta-llama/Llama-3.1-8B-Instruct",
-                    "meta-llama/Llama-3.1-70B-Instruct",
+                    "bigcomputer/jupycoder-7b-lora-350",
+                    "Qwen/Qwen2.5-Coder-7B-Instruct"
                 ],
                 label="Models"
             )

utils.py

@@ -317,4 +317,30 @@ def run_interactive_notebook(client, model, tokenizer, messages, sbx, max_new_to
             if tokens[-1] == "<|eot_id|>":
                 break
     
-    yield update_notebook_display(notebook_data), notebook_data, messages
+    yield update_notebook_display(notebook_data), notebook_data, messages
+
+def update_notebook_with_cell(notebook_data, code, output):
+    """Add a code cell and its output to the notebook"""
+    cell = {
+        "cell_type": "code",
+        "execution_count": None,
+        "metadata": {},
+        "source": code.split('\n'),
+        "outputs": [{
+            "output_type": "stream",
+            "name": "stdout",
+            "text": str(output).split('\n')
+        }] if output else []
+    }
+    notebook_data['cells'].append(cell)
+    return notebook_data
+
+def update_notebook_with_markdown(notebook_data, markdown_text):
+    """Add a markdown cell to the notebook"""
+    cell = {
+        "cell_type": "markdown",
+        "metadata": {},
+        "source": markdown_text.split('\n')
+    }
+    notebook_data['cells'].append(cell)
+    return notebook_data