Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoint/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+training_plot.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,118 @@
+---
+license: apache-2.0
+language:
+  - en
+library_name: transformers
+tags:
+  - smolagents
+  - code-generation
+  - qwen2
+  - text-generation
+pipeline_tag: text-generation
+base_model: Qwen/Qwen2.5-0.5B
+---
+
+# pynb-73m-base
+
+A 73M parameter language model trained for code generation with [smolagents](https://github.com/huggingface/smolagents). Built on the Qwen2 architecture.
+
+## Model Details
+
+| Property | Value |
+|----------|-------|
+| Parameters | 73.6M |
+| Architecture | Qwen2ForCausalLM |
+| Hidden size | 384 |
+| Layers | 12 |
+| Attention heads | 6 (2 KV heads, GQA 3:1) |
+| Intermediate size | 768 |
+| Context length | 2048 |
+| Vocab size | 151,671 |
+
+## Training
+
+Trained for 15,500 steps (~12 hours) on a single NVIDIA RTX 5070 Ti.
+
+![Training Progress](training_plot.png)
+
+| Metric | Start | End |
+|--------|-------|-----|
+| Train Loss | 287.8 | 53.4 |
+| Val Loss | 6.48 | 2.65 |
+
+## Quick Start with smolagents
+
+See [`inference_smolagent.py`](inference_smolagent.py) for full agent setup with LocalPythonExecutor and tools.
+
+```python
+from inference_smolagent import create_agent, CalculatorTool, FibonacciTool
+
+agent = create_agent(
+    model_id="AutomatedScientist/pynb-73m-base",
+    tools=[CalculatorTool(), FibonacciTool()],
+    max_steps=5,
+)
+
+result = agent.run("Calculate 15 * 7 + 23")
+print(result)
+```
+
+Or with HuggingFace API model:
+
+```python
+from smolagents import CodeAgent, HfApiModel
+
+model = HfApiModel(model_id="AutomatedScientist/pynb-73m-base")
+agent = CodeAgent(tools=[], model=model)
+
+result = agent.run("Calculate the sum of numbers from 1 to 100")
+print(result)
+```
+
+## Local Inference
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "AutomatedScientist/pynb-73m-base"  # or "checkpoint" for local
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+prompt = "Write a function to calculate fibonacci numbers"
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+outputs = model.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
+print(tokenizer.decode(outputs[0], skip_special_tokens=False))
+```
+
+## Inference Script
+
+See [`inference.py`](inference.py) for a wrapper class:
+
+```python
+from inference import CodeModel
+
+model = CodeModel("AutomatedScientist/pynb-73m-base")
+result = model.generate("Write a function to sort a list")
+print(result)
+```
+
+## Installation
+
+```bash
+pip install torch transformers smolagents
+```
+
+## Limitations
+
+- Small model (73M params) - limited reasoning capacity compared to larger models
+- Context window limited to 2,048 tokens
+- Best used with short prompts due to context constraints
+
+## License
+
+Apache 2.0

checkpoint/added_tokens.json

@@ -0,0 +1,30 @@
+{
+  "</tool_call>": 151658,
+  "<tool_call>": 151657,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|end_record_state|>": 151670,
+  "<|end_tool_call|>": 151666,
+  "<|end_tool_response|>": 151668,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|start_record_state|>": 151669,
+  "<|start_tool_call|>": 151665,
+  "<|start_tool_response|>": 151667,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

checkpoint/chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

checkpoint/config.json

@@ -0,0 +1,40 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "dtype": "float32",
+  "hidden_act": "silu",
+  "hidden_size": 384,
+  "initializer_range": 0.02,
+  "intermediate_size": 768,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 2048,
+  "max_window_layers": 28,
+  "model_type": "qwen2",
+  "num_attention_heads": 6,
+  "num_hidden_layers": 12,
+  "num_key_value_heads": 2,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 10000.0,
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "4.57.3",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151671
+}

checkpoint/generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.57.3",
+  "use_cache": false
+}

checkpoint/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c0f3750bf2a02e3450c09ce290922aecd910c7c43ecf18bac133c1b4f2228ae
+size 294393488

checkpoint/special_tokens_map.json

@@ -0,0 +1,60 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "<|start_tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|end_tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|start_tool_response|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|end_tool_response|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|start_record_state|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    },
+    {
+      "content": "<|end_record_state|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

checkpoint/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d508f3b5f4640a6623351b2a1e39389b41711492de2b19e6ad408461a89de0f
+size 11423080

checkpoint/tokenizer_config.json

@@ -0,0 +1,248 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<|start_tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151666": {
+      "content": "<|end_tool_call|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151667": {
+      "content": "<|start_tool_response|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151668": {
+      "content": "<|end_tool_response|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151669": {
+      "content": "<|start_record_state|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151670": {
+      "content": "<|end_record_state|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|start_tool_call|>",
+    "<|end_tool_call|>",
+    "<|start_tool_response|>",
+    "<|end_tool_response|>",
+    "<|start_record_state|>",
+    "<|end_record_state|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

inference.py

@@ -0,0 +1,73 @@
+"""inference.py - Code generation model wrapper for smolagents"""
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+class CodeModel:
+    def __init__(self, model_id: str, device: str = None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, fix_mistral_regex=True)
+        dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
+        self.model = AutoModelForCausalLM.from_pretrained(model_id).to(self.device, dtype=dtype)
+        self.model.eval()
+
+    def generate(self, prompt: str, max_new_tokens: int = 512, temperature: float = 0.7) -> str:
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                do_sample=True,
+                top_p=0.9,
+                repetition_penalty=1.2,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=self.tokenizer.eos_token_id,
+            )
+
+        new_tokens = outputs[0, inputs["input_ids"].shape[1]:]
+        return self.tokenizer.decode(new_tokens, skip_special_tokens=False)
+
+    def chat(self, messages: list[dict], max_new_tokens: int = 256) -> str:
+        """Generate response using chat template."""
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=False
+        )
+        inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.9,
+                repetition_penalty=1.2,
+            )
+
+        new_tokens = outputs[0, inputs["input_ids"].shape[1]:]
+        return self.tokenizer.decode(new_tokens, skip_special_tokens=False)
+
+
+if __name__ == "__main__":
+    import os
+    # Use local checkpoint if available, otherwise HuggingFace
+    model_id = "checkpoint" if os.path.exists("checkpoint") else "AutomatedScientist/pynb-73m-base"
+    model = CodeModel(model_id)
+
+    # Example: Generate code
+    result = model.generate("Write a Python function to calculate factorial")
+    print("Generated code:")
+    print(result)
+
+    # Example: Chat
+    messages = [
+        {"role": "system", "content": "You are a helpful coding assistant."},
+        {"role": "user", "content": "Write a function to reverse a string"}
+    ]
+    response = model.chat(messages)
+    print("\nChat response:")
+    print(response)

inference_smolagent.py

@@ -0,0 +1,347 @@
+"""inference_smolagent.py - Run model with smolagents CodeAgent and LocalPythonExecutor"""
+import os
+import re
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from smolagents import CodeAgent, Tool
+from smolagents.local_python_executor import LocalPythonExecutor
+from smolagents.models import ChatMessage, MessageRole, Model
+
+DEBUG = int(os.environ.get("DEBUG", 0))
+
+# Model's special tokens (from training)
+START_TOOL_CALL = "<|start_tool_call|>"
+END_TOOL_CALL = "<|end_tool_call|>"
+START_TOOL_RESPONSE = "<|start_tool_response|>"
+END_TOOL_RESPONSE = "<|end_tool_response|>"
+
+# Smolagents expected tokens
+SMOLAGENT_CODE_START = "<code>"
+SMOLAGENT_CODE_END = "</code>"
+
+
+class LocalCodeModel(Model):
+    """
+    Local model wrapper compatible with smolagents.
+
+    Handles translation between smolagents format and model's training format.
+    """
+
+    def __init__(self, model_id: str, device: str = None):
+        super().__init__()
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, fix_mistral_regex=True)
+        self.model = AutoModelForCausalLM.from_pretrained(model_id)
+        self.model.to(self.device)
+        self.model.eval()
+
+        # Cache special token IDs for stopping
+        self._end_tool_id = self.tokenizer.encode(END_TOOL_CALL, add_special_tokens=False)[-1]
+
+    def _convert_prompt_to_model_format(self, prompt: str) -> str:
+        """Convert smolagents prompt format to model's training format."""
+        # Replace smolagents code markers with model's markers
+        prompt = prompt.replace(SMOLAGENT_CODE_START, START_TOOL_CALL)
+        prompt = prompt.replace(SMOLAGENT_CODE_END, END_TOOL_CALL)
+        return prompt
+
+    def _convert_response_to_smolagent_format(self, response: str) -> str:
+        """Convert model's output format to smolagents expected format."""
+        # Replace model's markers with smolagents markers
+        response = response.replace(START_TOOL_CALL, SMOLAGENT_CODE_START)
+        response = response.replace(END_TOOL_CALL, SMOLAGENT_CODE_END)
+        response = response.replace(START_TOOL_RESPONSE, "")
+        response = response.replace(END_TOOL_RESPONSE, "")
+
+        # Clean up: remove orphan closing tags at start
+        response = re.sub(r'^\s*</code>\s*', '', response)
+
+        # Check if we have valid <code>...</code> block
+        has_open = SMOLAGENT_CODE_START in response
+        has_close = SMOLAGENT_CODE_END in response
+
+        # If only closing tag, remove it
+        if has_close and not has_open:
+            response = response.replace(SMOLAGENT_CODE_END, "")
+
+        # If no code markers, try to extract and wrap code
+        if SMOLAGENT_CODE_START not in response:
+            # Look for python code patterns in markdown
+            code_match = re.search(r'```(?:python)?\s*(.*?)\s*```', response, re.DOTALL)
+            if code_match:
+                code = code_match.group(1).strip()
+                if code:
+                    response = f"Thoughts: Executing the code\n{SMOLAGENT_CODE_START}\n{code}\n{SMOLAGENT_CODE_END}"
+            else:
+                # Look for any code-like content
+                lines = response.strip().split('\n')
+                code_lines = [l for l in lines if any(kw in l for kw in ['def ', 'print(', 'return ', '= ', 'import ', 'for ', 'if ', 'while '])]
+                if code_lines:
+                    code = '\n'.join(code_lines)
+                    response = f"Thoughts: Executing the code\n{SMOLAGENT_CODE_START}\n{code}\n{SMOLAGENT_CODE_END}"
+                else:
+                    # Fallback: wrap entire response as code if it looks like code
+                    clean = response.strip()
+                    if clean and not clean.startswith("Thoughts"):
+                        response = f"Thoughts: Attempting execution\n{SMOLAGENT_CODE_START}\nprint('No valid code generated')\n{SMOLAGENT_CODE_END}"
+
+        # Ensure closing tag exists if opening exists
+        if SMOLAGENT_CODE_START in response and SMOLAGENT_CODE_END not in response:
+            response = response + f"\n{SMOLAGENT_CODE_END}"
+
+        return response
+
+    def generate(
+        self,
+        messages: list[ChatMessage],
+        stop_sequences: list[str] | None = None,
+        grammar: str | None = None,
+        tools_to_call_from: list[Tool] | None = None,
+        **kwargs,
+    ) -> ChatMessage:
+        """Generate response for message history (required by smolagents Model)."""
+        # Debug: show what messages are passed (including executor output)
+        if DEBUG:
+            print("\n[DEBUG] Messages received by model:")
+            for i, msg in enumerate(messages):
+                role = msg.role.value if hasattr(msg.role, "value") else msg.role
+                content = str(msg.content)[:200] if msg.content else "<empty>"
+                print(f"  [{i}] {role}: {content}...")
+            print()
+
+        # Convert ChatMessage objects to dicts for chat template
+        messages_dicts = []
+        for msg in messages:
+            if hasattr(msg, "role") and hasattr(msg, "content"):
+                role = msg.role.value if hasattr(msg.role, "value") else str(msg.role)
+                content = msg.content if isinstance(msg.content, str) else str(msg.content or "")
+                # Convert prompt format in content
+                content = self._convert_prompt_to_model_format(content)
+                # Wrap observations (executor output) in tool response tokens
+                if "Observation:" in content or "Out:" in content:
+                    # Extract the observation content
+                    obs_match = re.search(r'(?:Observation:|Out:)\s*(.*)', content, re.DOTALL)
+                    if obs_match:
+                        obs_content = obs_match.group(1).strip()
+                        content = f"{START_TOOL_RESPONSE}\n{obs_content}\n{END_TOOL_RESPONSE}"
+                messages_dicts.append({"role": role, "content": content})
+            else:
+                messages_dicts.append(msg)
+
+        # Convert messages to prompt using chat template
+        prompt = self.tokenizer.apply_chat_template(
+            messages_dicts,
+            add_generation_prompt=True,
+            tokenize=False
+        )
+
+        # Check prompt length
+        if DEBUG:
+            full_tokens = self.tokenizer(prompt, return_tensors="pt")
+            print(f"[DEBUG] Prompt length: {full_tokens['input_ids'].shape[1]} tokens (max: 2048)")
+
+        # Truncate to fit model's context window (2048 tokens, leave room for generation)
+        max_input_tokens = 1536  # Leave 512 for generation
+        inputs = self.tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=max_input_tokens
+        ).to(self.device)
+
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=512,
+                temperature=0.7,
+                do_sample=True,
+                top_p=0.9,
+                repetition_penalty=1.2,
+                pad_token_id=self.tokenizer.pad_token_id,
+                eos_token_id=[self.tokenizer.eos_token_id, self._end_tool_id],
+            )
+
+        new_tokens = outputs[0, inputs["input_ids"].shape[1]:]
+        response = self.tokenizer.decode(new_tokens, skip_special_tokens=False)
+
+        # Handle stop sequences
+        if stop_sequences:
+            for seq in stop_sequences:
+                if seq in response:
+                    response = response.split(seq)[0]
+
+        # Convert response format for smolagents
+        response = self._convert_response_to_smolagent_format(response)
+
+        return ChatMessage(role=MessageRole.ASSISTANT, content=response)
+
+
+# Example tools
+class CalculatorTool(Tool):
+    name = "calculator"
+    description = "Evaluates a mathematical expression and returns the result."
+    inputs = {
+        "expression": {
+            "type": "string",
+            "description": "The mathematical expression to evaluate (e.g., '2 + 2 * 3')"
+        }
+    }
+    output_type = "number"
+
+    def forward(self, expression: str) -> float:
+        # Safe eval for math expressions
+        allowed = set("0123456789+-*/().^ ")
+        if not all(c in allowed for c in expression):
+            raise ValueError("Invalid characters in expression")
+        return eval(expression.replace("^", "**"))
+
+
+class FibonacciTool(Tool):
+    name = "fibonacci"
+    description = "Calculate the nth Fibonacci number."
+    inputs = {
+        "n": {
+            "type": "integer",
+            "description": "The position in Fibonacci sequence (0-indexed)"
+        }
+    }
+    output_type = "integer"
+
+    def forward(self, n: int) -> int:
+        if n < 0:
+            raise ValueError("n must be non-negative")
+        if n <= 1:
+            return n
+        a, b = 0, 1
+        for _ in range(2, n + 1):
+            a, b = b, a + b
+        return b
+
+
+SHORT_PROMPT_TEMPLATES = {
+    "system_prompt": """You solve tasks by writing Python code.
+
+Rules:
+- Write code inside <code> and </code> tags
+- Use print() to show results
+- Use final_answer(result) when done
+
+Format:
+Thoughts: your reasoning
+<code>
+# your code
+</code>""",
+    "planning": {
+        "initial_plan": "",
+        "update_plan_pre_messages": "",
+        "update_plan_post_messages": "",
+    },
+    "managed_agent": {
+        "task": "",
+        "report": "",
+    },
+    "final_answer": {
+        "pre_messages": "",
+        "post_messages": "",
+    },
+}
+
+
+def create_agent(
+    model_id: str = "AutomatedScientist/pynb-73m-base",
+    tools: list[Tool] | None = None,
+    additional_authorized_imports: list[str] | None = None,
+    max_steps: int = 5,
+    use_short_prompt: bool = True,
+) -> CodeAgent:
+    """
+    Create a CodeAgent with LocalPythonExecutor.
+
+    Args:
+        model_id: HuggingFace model ID or local path
+        tools: List of tools to provide to the agent
+        additional_authorized_imports: Extra imports to allow in executor
+        max_steps: Maximum agent steps before stopping
+        use_short_prompt: Use shorter system prompt for small context models
+
+    Returns:
+        Configured CodeAgent instance
+    """
+    model = LocalCodeModel(model_id)
+
+    # Default authorized imports
+    authorized_imports = [
+        "math", "statistics", "random", "datetime",
+        "collections", "itertools", "re", "json",
+        "functools", "operator"
+    ]
+    if additional_authorized_imports:
+        authorized_imports.extend(additional_authorized_imports)
+
+    # Create executor with sandbox
+    executor = LocalPythonExecutor(
+        additional_authorized_imports=authorized_imports,
+        max_print_outputs_length=10000,
+    )
+
+    # Build agent config
+    agent_kwargs = {
+        "tools": tools or [],
+        "model": model,
+        "executor": executor,
+        "max_steps": max_steps,
+        "verbosity_level": 1,
+    }
+
+    # Use short prompt for small context models
+    if use_short_prompt:
+        agent_kwargs["prompt_templates"] = SHORT_PROMPT_TEMPLATES
+
+    agent = CodeAgent(**agent_kwargs)
+
+    return agent
+
+
+def run_task(agent: CodeAgent, task: str) -> any:
+    """
+    Run a task through the agent.
+
+    Args:
+        agent: CodeAgent instance
+        task: Natural language task description
+
+    Returns:
+        Agent output
+    """
+    print(f"\n{'='*60}")
+    print(f"Task: {task}")
+    print(f"{'='*60}\n")
+
+    result = agent.run(task)
+
+    print(f"\n{'='*60}")
+    print(f"Result: {result}")
+    print(f"{'='*60}\n")
+
+    return result
+
+
+if __name__ == "__main__":
+    import sys
+
+    # Use local checkpoint if available, otherwise HuggingFace
+    model_id = "checkpoint" if os.path.exists("checkpoint") else "AutomatedScientist/pynb-73m-base"
+
+    agent = create_agent(
+        model_id=model_id,
+        tools=[CalculatorTool(), FibonacciTool()],
+        max_steps=8,
+    )
+
+    # Run example task
+    task = sys.argv[1] if len(sys.argv) > 1 else "Calculate 15 * 7 + 23"
+    try:
+        result = run_task(agent, task)
+    except Exception as e:
+        print(f"Error: {e}")