README.md

---
language:
- en
license: apache-2.0
tags:
- text-to-code
- manim
- python
- mistral
- fine-tuned
- lora
- qlora
base_model: mistralai/Mistral-7B-v0.3
datasets:
- Edoh/manim_python
metrics:
- accuracy
pipeline_tag: text-generation
library_name: transformers
---

# Mistral Manim Python Coder (`TheSon2202/mistral-manim-python-coder-v01`)

This model is a fine-tuned version of **Mistral-7B-v0.3** using **QLoRA (4-bit NF4)**, specialized in translating natural language instructions (**Text-to-Instruction**) into precise Python code for the mathematical animation library **Manim**.

---

## 1. Hyperparameters & Configuration

| Configuration Parameter | Value |
| :--- | :--- |
| **Base Model** | `mistralai/Mistral-7B-v0.3` |
| **Dataset** | `Edoh/manim_python` |
| **Maximum Sequence Length** | `512` tokens |
| **Learning Rate** | `2e-4` (0.0002) |
| **Weight Decay** | `0.03` |
| **Per-Device Batch Size** | `2` |
| **Gradient Accumulation Steps** | `4` |
| **Number of Epochs** | `2` (Total 120 steps) |
| **Optimizer** | `paged_adamw_32bit` |
| **LR Scheduler** | `cosine` |
| **Gradient Clipping (`max_grad_norm`)** | `0.3` |
| **Warmup Steps Ratio** | `0.1` (10%) |

### PEFT (LoRA) Config
* **Rank (`r`):** `16`
* **Alpha (`lora_alpha`):** `32`
* **Dropout (`lora_dropout`):** `0.05`
* **Target Modules:** `["q_proj", "k_proj", "v_proj", "o_proj"]`
* **Task Type:** `CAUSAL_LM`

### Quantization Config (BitsAndBytes)
* **Load in 4-bit:** `True`
* **Quant Type:** `nf4` (Normal Float 4)
* **Compute Dtype:** `torch.float16`
* **Double Quantization:** `True`

---

## 2. Training Metrics & Evaluation Results

The training process recorded convergence milestones across checkpoints (saved periodically every 50 steps):

| Training Step | Training Loss | Validation Loss | Num Tokens | Mean Token Accuracy |
| :---: | :---: | :---: | :---: | :---: |
| **Step 50** | `0.2506` | `0.2504` | 41,922 | **94.41%** |
| **Step 100** | `0.2271` | `0.2374` | 83,632 | **94.83%** |
| **Step 120 (Final)** | `0.2259` | `0.2359` | 100,332 | **94.88%** |


![Screenshot 2026-08-05 at 03.12.42](https://cdn-uploads.huggingface.co/production/uploads/68c7b25941cddd7ea63d8e89/9HeDsCgzfa3dBj2oSgGWk.png)

> **General Overview:** Both training and validation losses decreased steadily and closely tracked each other (showing no signs of overfitting). Combined with an average token accuracy of approximately **94.88%**, this demonstrates that the model successfully learned Manim's syntax and programming conventions.

---

## 3. Inference Demo

You can load the model directly from the Hugging Face Hub to generate Manim code using the following Python snippet:

```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "TheSon2202/mistral-manim-python-coder-v01"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.float16
)

# Configure Chat Template for Mistral Base Model
tokenizer.chat_template = (
    "{{ bos_token }}"
    "{% for message in messages %}"
        "{% if message['role'] == 'system' %}"
            "{{ 'System: ' + message['content'] + '\n\n' }}"
        "{% elif message['role'] == 'user' %}"
            "{{ '[INST] ' + message['content'] + ' [/INST]' }}"
        "{% elif message['role'] == 'assistant' %}"
            "{{ ' ' + message['content'] + eos_token }}"
        "{% endif %}"
    "{% endfor %}"
)

def generate_manim_code(instruction):
    system_prompt = "Yor are an Coding Python Expert, read the instruction and complete these code correctly"
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": instruction}
    ]
    
    prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.2,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Test code generation
test_instruction = "Create a square with side length 4 and color it red, then animate it to shift right by 3 units."
print(generate_manim_code(test_instruction))
```
---

### 📤 Expected Output (Clean Python Code)

```python
from manim import *

class MyScene(Scene):
    def construct(self):
        square = Square(side_length=4, color=RED)
        self.add(square)
        self.play(square.animate.shift(RIGHT * 3), run_time=3)
```