.

app.py

@@ -5,6 +5,8 @@ from e2b_code_interpreter import Sandbox
 from pathlib import Path
 from peft import PeftModel
 from transformers import AutoTokenizer,AutoModelForCausalLM
+from huggingface_hub import snapshot_download
+from vllm import LLM, SamplingParams
 import json
 
 if not get_space():
@@ -20,6 +22,7 @@ from utils import (
     run_interactive_notebook,
     create_base_notebook,
     update_notebook_display,
+    user_template,
 )
 
 E2B_API_KEY = os.environ["E2B_API_KEY"]
@@ -48,17 +51,12 @@ def execute_jupyter_agent(
     os.makedirs(save_dir, exist_ok=True)
     save_dir = os.path.join(save_dir, 'jupyter-agent.ipynb')
 
-    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct")
-    model = AutoModelForCausalLM.from_pretrained(
-        "Qwen/Qwen2.5-Coder-7B-Instruct", torch_dtype='auto'
-    ).eval()
-    # # Load the LoRA adapter and move the model to GPU
-    model = PeftModel.from_pretrained(
-        model,
-        model_name,
-        device_map="auto",  # Automatically allocate model layers to available devices
-        trust_remote_code=True
-    ).eval()
+    sampling_params = SamplingParams(
+        temperature=0.2,
+        max_tokens=512,
+    )
+
+    lora_path = snapshot_download(model_name)
 
     filenames = []
     if files is not None:
@@ -69,7 +67,7 @@ def execute_jupyter_agent(
                 sbx.files.write(filpath.name, file)
                 filenames.append(filpath.name)
 
-    # Initialize message_history if it doesn't exist
+    # Initialize message_history and notebook_data if they don't exist
     if len(message_history) == 0:
         message_history.append(
             {
@@ -77,24 +75,49 @@ def execute_jupyter_agent(
                 "content": system_prompt.format("- " + "\n- ".join(filenames)),
             }
         )
-    message_history.append({"role": "user", "content": user_input})
+        current_notebook = None
+    else:
+        # Load existing notebook data from file
+        try:
+            with open(save_dir, 'r', encoding='utf-8') as f:
+                current_notebook = json.load(f)
+        except (FileNotFoundError, json.JSONDecodeError):
+            current_notebook = None
+
+    # Add user input with is_user_prompt flag
+    message_history.append({"role": "user", "content": user_input, "is_user_prompt": True})
 
     print("history:", message_history)
 
+    # Update notebook with new user prompt if we have an existing notebook
+    if current_notebook is not None:
+        current_notebook["cells"].append({
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": user_template.format(user_input.replace('\n', '<br>'))
+        })
+        # Save the updated notebook
+        with open(save_dir, 'w', encoding='utf-8') as f:
+            json.dump(current_notebook, f, indent=2)
+
     for notebook_html, notebook_data, messages in run_interactive_notebook(
-        model, tokenizer, message_history, sbx, max_new_tokens=max_new_tokens
+        lora_path, sampling_params, message_history, sbx, notebook_data=current_notebook, max_new_tokens=max_new_tokens
     ):
         message_history = messages
         
-        yield notebook_html, message_history, TMP_DIR+"jupyter-agent.ipynb"
-    
-    with open(save_dir, 'w', encoding='utf-8') as f:
-        json.dump(notebook_data, f, indent=2)
-    yield notebook_html, message_history, save_dir
+        # Save notebook after each update
+        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 = []
-    return update_notebook_display(create_base_notebook([])[0]), msg_state
+    # Also clear the notebook file
+    notebook_data = create_base_notebook([])[0]
+    with open(TMP_DIR+"jupyter-agent.ipynb", 'w', encoding='utf-8') as f:
+        json.dump(notebook_data, f, indent=2)
+    return update_notebook_display(notebook_data), msg_state
 
 
 css = """
@@ -151,9 +174,9 @@ with gr.Blocks() as demo:
             )
 
             model = gr.Dropdown(
-                value="bigcomputer/jupycoder-7b-lora-350",
+                value="bigcomputer/jupycoder-7b-lora-200",
                 choices=[
-                    "bigcomputer/jupycoder-7b-lora-350",
+                    "bigcomputer/jupycoder-7b-lora-200",
                     "Qwen/Qwen2.5-Coder-7B-Instruct"
                 ],
                 label="Models"

requirements.txt

@@ -4,4 +4,4 @@ huggingface_hub
 e2b-code-interpreter
 transformers
 traitlets
-peft
+vllm

utils.py

@@ -3,12 +3,14 @@ from nbformat.v4 import new_notebook, new_markdown_cell, new_code_cell
 from nbconvert import HTMLExporter
 from huggingface_hub import InferenceClient
 from e2b_code_interpreter import Sandbox
-from transformers import AutoTokenizer
+from vllm.lora.request import LoRARequest
 from traitlets.config import Config
+from vllm import LLM
 import re
 
 config = Config()
 html_exporter = HTMLExporter(config=config, template_name="classic")
+BASE_MODEL = LLM(model="Qwen/Qwen2.5-Coder-7B-Instruct", enable_lora=True)
 
 # Constants
 MAX_TURNS = 10
@@ -157,11 +159,20 @@ def create_base_notebook(messages):
                 "source": text
                 })
         elif message["role"] == "user":
-            text = user_template.format(message["content"].replace('\n', '<br>'))
-            base_notebook["cells"].append({
-                "cell_type": "markdown",
-                "metadata": {},
-                "source": text
+            # Check if this is an actual user prompt (has is_user_prompt flag)
+            if message.get("is_user_prompt", False):
+                text = user_template.format(message["content"].replace('\n', '<br>'))
+                base_notebook["cells"].append({
+                    "cell_type": "markdown",
+                    "metadata": {},
+                    "source": text
+                    })
+            else:
+                # This is an execution output, add as code cell output
+                base_notebook["cells"][-1]["outputs"].append({
+                    "output_type": "stream",
+                    "name": "stdout",
+                    "text": message["content"]
                 })
 
         elif message["role"] == "assistant" and "tool_calls" in message:
@@ -219,84 +230,144 @@ def update_notebook_display(notebook_data):
     notebook_body = notebook_body.replace(bad_html_bad, "")
     return notebook_body
 
-def run_interactive_notebook(model, tokenizer, messages, sbx, max_new_tokens=512):
-    notebook_data, code_cell_counter = create_base_notebook(messages)
-    turns = 0
+def run_interactive_notebook(lora_path, sampling_params, messages, sbx, notebook_data=None, max_new_tokens=512):
+    """
+    Run interactive notebook with model.
+    
+    Args:
+        lora_path: Path to LoRA adapter
+        sampling_params: Sampling parameters for the model
+        messages: List of conversation messages
+        sbx: Sandbox environment for code execution
+        notebook_data: Existing notebook data when continuing a session
+        max_new_tokens: Maximum number of new tokens to generate
+    """
+    # For first run or when notebook_data is not provided
+    if notebook_data is None:
+        # Create a separate list for display messages with is_user_prompt flag
+        display_messages = []
+        model_messages = []  # Clean messages for model
+        for msg in messages:
+            display_msg = msg.copy()
+            if msg["role"] == "user":
+                display_msg["is_user_prompt"] = True
+            display_messages.append(display_msg)
+            model_messages.append(msg.copy())  # Keep clean copy for model
+        notebook_data, code_cell_counter = create_base_notebook(display_messages)
+    else:
+        # For subsequent runs, use existing messages but clean them for model
+        display_messages = messages
+        model_messages = []
+        for msg in messages:
+            # Create clean copy without display flags for model
+            model_msg = msg.copy()
+            if "is_user_prompt" in model_msg:
+                del model_msg["is_user_prompt"]
+            model_messages.append(model_msg)
+            
+        # Find the last code cell counter
+        code_cell_counter = 0
+        for cell in notebook_data["cells"]:
+            if cell["cell_type"] == "code" and cell.get("execution_count"):
+                code_cell_counter = max(code_cell_counter, cell["execution_count"])
     
-    while turns <= MAX_TURNS:
+    turns = 0
+    while turns < MAX_TURNS:
         turns += 1
-        # Generate response using the model
-        text = tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-
-        generated_ids = model.generate(
-            **model_inputs,
-            max_new_tokens=max_new_tokens
+        # Generate response using the model with clean messages
+        print(model_messages)
+        response_stream = BASE_MODEL.chat(
+            model_messages,
+            sampling_params,
+            lora_request=LoRARequest("lora_adapter", 1, lora_path),
+            add_generation_prompt=True
+        )[0].outputs[0].text
+        
+        # Check for duplicate responses
+        is_duplicate = any(
+            msg["role"] == "assistant" and msg["content"].strip() == response_stream.strip()
+            for msg in model_messages
         )
-        generated_ids = [
-            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-        ]
-        response_stream = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
         
-        # Process the full response at once
-        parts = re.split(r'(```python[\s\S]*?```)', response_stream)
+        if is_duplicate:
+            # If duplicate found, yield current state and break
+            yield update_notebook_display(notebook_data), notebook_data, display_messages
+            break
         
-        for part in parts:
-            if part.strip():
-                if part.startswith('```python'):
-                    # Extract code without the markers
-                    code = re.sub(r'```python\n|```', '', part).strip()
-                    code_cell_counter += 1
-                    
-                    # Add code cell
-                    notebook_data["cells"].append({
-                        "cell_type": "code",
-                        "execution_count": code_cell_counter,
-                        "metadata": {},
-                        "source": code,
-                        "outputs": []
-                    })
-                    
-                    # Execute code
-                    exec_result, execution = execute_code(sbx, code)
-                    messages.append({
-                        "role": "assistant",
-                        "content": code,
-                        "tool_calls": [{
-                            "type": "function",
-                            "function": {
-                                "name": "code_interpreter",
-                                "arguments": {"code": code}
-                            }
-                        }]
-                    })
-                    messages.append({
-                        "role": "ipython", 
-                        "content": parse_exec_result_llm(execution), 
-                        "nbformat": parse_exec_result_nb(execution)
-                    })
-                    
-                    # Update cell with execution results
-                    notebook_data["cells"][-1]["outputs"] = parse_exec_result_nb(execution)
-                else:
-                    # Add markdown cell for non-code content
-                    notebook_data["cells"].append({
-                        "cell_type": "markdown",
-                        "metadata": {},
-                        "source": part.strip()
-                    })
-                    messages.append({
-                        "role": "assistant",
-                        "content": part.strip()
-                    })
+        # Add the full response as an assistant message
+        assistant_msg = {
+            "role": "assistant",
+            "content": response_stream
+        }
+        model_messages.append(assistant_msg.copy())
+        display_messages.append(assistant_msg)
         
-        # Return the final result
-        yield update_notebook_display(notebook_data), notebook_data, messages
-        break
+        # Check if response contains code block
+        code_match = re.search(r'```python\n(.*?)```', response_stream, re.DOTALL)
+        if code_match:
+            # Extract and execute the code
+            code = code_match.group(1).strip()
+            code_cell_counter += 1
+            
+            # Add code cell
+            notebook_data["cells"].append({
+                "cell_type": "code",
+                "execution_count": code_cell_counter,
+                "metadata": {},
+                "source": code,
+                "outputs": []
+            })
+            
+            # Execute code and get results
+            exec_result, execution = execute_code(sbx, code)
+            
+            # Get execution results in notebook format
+            outputs = parse_exec_result_nb(execution)
+            
+            # Create text-only version for user message
+            user_content = []
+            for output in outputs:
+                if output.get('output_type') == 'stream':
+                    user_content.append(output['text'])
+                elif output.get('output_type') == 'error':
+                    user_content.append('\n'.join(output['traceback']))
+                elif output.get('output_type') in ['execute_result', 'display_data']:
+                    data = output.get('data', {})
+                    if 'text/plain' in data:
+                        user_content.append('\n'.join(data['text/plain']))
+                    if any(key.startswith('image/') for key in data.keys()):
+                        user_content.append('<image>')
+            
+            # Create execution result message
+            user_msg = {
+                "role": "user", 
+                "content": '\n'.join(user_content)
+            }
+            # Add clean version to model messages
+            model_messages.append(user_msg.copy())
+            # Add version with display flag to display messages
+            display_msg = user_msg.copy()
+            display_msg["is_user_prompt"] = False
+            display_messages.append(display_msg)
+            
+            # Update cell with execution results
+            notebook_data["cells"][-1]["outputs"] = outputs
+            
+            # Yield intermediate results after each turn
+            yield update_notebook_display(notebook_data), notebook_data, display_messages
+        else:
+            # No code in this turn, add as markdown and break
+            notebook_data["cells"].append({
+                "cell_type": "markdown",
+                "metadata": {},
+                "source": response_stream
+            })
+            # Yield final results and break
+            yield update_notebook_display(notebook_data), notebook_data, display_messages
+            break
     
-    yield update_notebook_display(notebook_data), notebook_data, messages
+    # Final yield in case we hit MAX_TURNS
+    yield update_notebook_display(notebook_data), notebook_data, display_messages
 
 def update_notebook_with_cell(notebook_data, code, output):
     """Add a code cell and its output to the notebook"""