update readme

README.md

@@ -1,199 +1,261 @@
 ---
+license: apache-2.0
 library_name: transformers
-tags: []
+tags:
+- code
+- jupyter
+- agent
+- data-science
+- qwen
+- thinking
+base_model: Qwen/Qwen3-4B-Thinking-2507
+datasets:
+- data-agents/jupyter-agent-dataset
+language:
+- en
+- code
+pipeline_tag: text-generation
 ---
 
-# Model Card for Model ID
+# Jupyter Agent Qwen3-4B Thinking
 
-<!-- Provide a quick summary of what the model is/does. -->
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/650ed7adf141bc34f91a12ae/ZyF9foqe5SLECwkq0dOpT.png)
 
+**Jupyter Agent Qwen3-4B Thinking** is a fine-tuned version of [Qwen3-4B-Thinking-2507](https://huggingface.co/Qwen/Qwen3-4B-Thinking-2507) specifically optimized for **data science agentic tasks** in Jupyter notebook environments. This model can execute Python code, analyze datasets, and provide step-by-step reasoning with intermediate computations to solve realistic data analysis problems.
 
+- **Model type:** Causal Language Model (Thinking)
+- **Language(s):** English, Python
+- **License:** Apache 2.0
+- **Finetuned from:** [Qwen/Qwen3-4B-Thinking-2507](https://huggingface.co/Qwen/Qwen3-4B-Thinking-2507)
 
-## Model Details
+## Key Features
 
-### Model Description
-
-<!-- Provide a longer summary of what this model is. -->
-
-This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
-
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
-
-### Model Sources [optional]
-
-<!-- Provide the basic links for the model. -->
-
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
-
-## Uses
-
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
-
-### Direct Use
-
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
-
-[More Information Needed]
-
-### Downstream Use [optional]
-
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
-
-[More Information Needed]
-
-### Out-of-Scope Use
-
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
-
-[More Information Needed]
-
-## Bias, Risks, and Limitations
-
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
-
-[More Information Needed]
-
-### Recommendations
-
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
-
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
-
-## How to Get Started with the Model
-
-Use the code below to get started with the model.
-
-[More Information Needed]
+- **Jupyter-native agent** that lives inside notebook environments
+- **Code execution** with pandas, numpy, matplotlib, and other data science libraries
+- **Step-by-step reasoning** with intermediate computations and thinking traces
+- **Dataset-grounded analysis** trained on real Kaggle notebook workflows
+- **Tool calling** for structured code execution and final answer generation
+
+## Performance
+
+On the [DABStep benchmark](https://huggingface.co/spaces/adyen/DABstep) for data science tasks:
+
+| Model | Easy Tasks | Hard Tasks |
+|-------|------------|------------|
+| Qwen3-4B-Thinking-2507 (Base) | 44.0% | 2.1% |
+| **Jupyter Agent Qwen3-4B Thinking** | **70.8%** | **3.4%** |
+
+**State-of-the-art performance** for small models on realistic data analysis tasks.
+
+## Model Sources
+
+- **Repository:** [jupyter-agent](https://github.com/huggingface/jupyter-agent)
+- **Dataset:** [jupyter-agent-dataset](https://huggingface.co/datasets/data-agents/jupyter-agent-dataset)
+- **Blog post:** [Jupyter Agents: training LLMs to reason with notebooks](https://huggingface.co/blog/jupyter-agent-2)
+- **Demo:** [Jupyter Agent 2](https://huggingface.co/spaces/lvwerra/jupyter-agent-2)
+
+## Usage
+
+### Basic Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "data-agents/jupyter-agent-qwen3-4b-thinking"
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    device_map="auto"
+)
+
+# Prepare input
+prompt = "Analyze this sales dataset and find the top 3 performing products by revenue."
+messages = [
+    {"role": "user", "content": prompt}
+]
+
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True
+)
+model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+# Generate response
+generated_ids = model.generate(
+    **model_inputs,
+    max_new_tokens=16384
+)
+output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+```
+
+### Decoding Thinking and Content
+
+For thinking models, you can extract both the reasoning and final response:
+
+```python
+try:
+    # Find the end of thinking section (</think>)
+    index = len(output_ids) - output_ids[::-1].index(151668)
+except ValueError:
+    index = 0
+
+thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
+content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+
+print("Thinking:", thinking_content)
+print("Response:", content)
+```
+
+### Agentic Usage with Tool Calling
+
+The model works best with proper scaffolding for tool calling:
+
+```python
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "execute_code",
+            "description": "Execute Python code in a Jupyter environment",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "code": {
+                        "type": "string",
+                        "description": "Python code to execute"
+                    }
+                },
+                "required": ["code"]
+            }
+        }
+    },
+    {
+        "type": "function", 
+        "function": {
+            "name": "final_answer",
+            "description": "Provide the final answer to the question",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "answer": {
+                        "type": "string",
+                        "description": "The final answer"
+                    }
+                },
+                "required": ["answer"]
+            }
+        }
+    }
+]
+
+# Include tools in the conversation
+messages = [
+    {
+        "role": "system", 
+        "content": "You are a data science assistant. Use the available tools to analyze data and provide insights."
+    },
+    {"role": "user", "content": prompt}
+]
+```
 
 ## Training Details
 
 ### Training Data
 
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+The model was fine-tuned on the [Jupyter Agent Dataset](https://huggingface.co/datasets/data-agents/jupyter-agent-dataset), which contains:
 
-[More Information Needed]
+- **51,389 synthetic notebooks** (~2B tokens)
+- **Dataset-grounded QA pairs** from real Kaggle notebooks
+- **Executable reasoning traces** with intermediate computations
+- **High-quality educational content** filtered and scored by LLMs
 
 ### Training Procedure
 
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
-
-#### Preprocessing [optional]
-
-[More Information Needed]
-
+- **Base Model:** Qwen3-4B-Thinking-2507
+- **Training Method:** Full-parameter fine-tuning (not PEFT)
+- **Optimizer:** AdamW with cosine learning rate scheduling
+- **Learning Rate:** 5e-6
+- **Epochs:** 5 (optimal based on ablation study)
+- **Context Length:** 32,768 tokens
+- **Batch Size:** Distributed across multiple GPUs
+- **Loss:** Assistant-only loss (`assistant_loss_only=True`)
+- **Regularization:** NEFTune noise (α=7) for full-parameter training
 
-#### Training Hyperparameters
+### Training Infrastructure
 
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
-
-#### Speeds, Sizes, Times [optional]
-
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
-
-[More Information Needed]
+- **Framework:** [TRL](https://github.com/huggingface/trl) with [Transformers](https://github.com/huggingface/transformers)
+- **Distributed Training:** DeepSpeed ZeRO-2 across multiple nodes
+- **Hardware:** Multi-GPU setup with SLURM orchestration
 
 ## Evaluation
 
-<!-- This section describes the evaluation protocols and provides the results. -->
-
-### Testing Data, Factors & Metrics
-
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
-
-[More Information Needed]
-
-#### Factors
-
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
-[More Information Needed]
-
-#### Metrics
-
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
-[More Information Needed]
-
-### Results
-
-[More Information Needed]
+### Benchmark: DABStep
 
-#### Summary
+The model was evaluated on [DABStep](https://huggingface.co/spaces/adyen/DABstep), a benchmark for data science agents with realistic tasks involving:
 
+- **Dataset analysis** with pandas and numpy
+- **Visualization** with matplotlib/seaborn
+- **Statistical analysis** and business insights
+- **Multi-step reasoning** with intermediate computations
 
+The model achieves **26.8% improvement** over the base model and **11.1% improvement** over scaffolding alone.
 
-## Model Examination [optional]
+<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/jupyter-agent-2/training_dabstep_easy.png" alt="DABstep Easy Score"/>
 
-<!-- Relevant interpretability work for the model goes here -->
+We can also see, that the hard score can increase too even though our dataset is focused on easier questions.
 
-[More Information Needed]
+<img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/jupyter-agent-2/training_dabstep_hard.png" alt="DABstep Hard Score"/>
 
-## Environmental Impact
+## Limitations and Bias
 
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+### Technical Limitations
 
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+- **Context window:** Limited to 32K tokens, may struggle with very large notebooks
+- **Tool calling format:** Requires specific scaffolding for optimal performance  
+- **Dataset domains:** Primarily trained on Kaggle-style data science tasks
+- **Code execution:** Requires proper sandboxing for safe execution
 
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
+### Potential Biases
 
-## Technical Specifications [optional]
+- **Domain bias:** Trained primarily on Kaggle notebooks, may not generalize to all data science workflows
+- **Language bias:** Optimized for English and Python, limited multilingual support
+- **Task bias:** Focused on structured data analysis, may underperform on unstructured data tasks
 
-### Model Architecture and Objective
-
-[More Information Needed]
-
-### Compute Infrastructure
-
-[More Information Needed]
-
-#### Hardware
-
-[More Information Needed]
-
-#### Software
-
-[More Information Needed]
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
-
-[More Information Needed]
-
-**APA:**
-
-[More Information Needed]
-
-## Glossary [optional]
+### Recommendations
 
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+- Use in **sandboxed environments** like [E2B](https://e2b.dev/) for safe code execution
+- **Validate outputs** before using in production systems
+- **Review generated code** for security and correctness
+- Consider **domain adaptation** for specialized use cases
 
-[More Information Needed]
+## Ethical Considerations
 
-## More Information [optional]
+- **Code Safety:** Always execute generated code in secure, isolated environments
+- **Data Privacy:** Be cautious when analyzing sensitive datasets
+- **Verification:** Validate all analytical conclusions and insights
+- **Attribution:** Acknowledge model assistance in data analysis workflows
 
-[More Information Needed]
+## Citation
 
-## Model Card Authors [optional]
+```bibtex
+@misc{jupyteragentqwen3thinking,
+  title={Jupyter Agent Qwen3-4B Thinking},
+  author={Baptiste Colle and Hanna Yukhymenko and Leandro von Werra},
+  year={2025},
+  publisher={Hugging Face},
+  url={https://huggingface.co/data-agents/jupyter-agent-qwen3-4b-thinking}
+}
+```
 
-[More Information Needed]
+## Related Work
 
-## Model Card Contact
+- **Dataset:** [jupyter-agent-dataset](https://huggingface.co/datasets/data-agents/jupyter-agent-dataset)
+- **Non-thinking version:** [jupyter-agent-qwen3-4b-instruct](https://huggingface.co/data-agents/jupyter-agent-qwen3-4b-instruct)
+- **Base model:** [Qwen3-4B-Thinking-2507](https://huggingface.co/Qwen/Qwen3-4B-Thinking-2507)
+- **Benchmark:** [DABStep](https://huggingface.co/spaces/adyen/DABstep)
 
-[More Information Needed]
+*For more details, see our [blog post](https://huggingface.co/blog/jupyter-agent-2) and [GitHub repository](https://github.com/huggingface/jupyter-agent).*