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%

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:

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)

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)