README.md

---
base_model: unsloth/qwen2.5-coder-7b-bnb-4bit
tags:
- text-generation-inference
- transformers
- unsloth
- qwen2
- diagram
- text-to-diagram
license: apache-2.0
language:
- en
datasets:
- huytd189/pintora-instruct
- huytd189/pintora-edit-instruct
---

# Pintora-Coder-7B

## Introduction

Pintora-Coder-7B is a fine-tuned version of [Qwen2.5-Coder-7B](https://huggingface.co/Qwen/Qwen2.5-Coder-7B-Instruct) to support the [Pintora](https://github.com/hikerpig/pintora) diagram language.

The model supports the following features:

1. Generate diagrams from scratch.
2. Edit existing diagrams.

## Training Details

The model has been trained in the following steps:

1. Continued pretraining with the [pintora-instruct](https://huggingface.co/datasets/huytd189/pintora-instruct) dataset.
2. Instruction fine-tuned with the [pintora-edit-instruct](https://huggingface.co/datasets/huytd189/pintora-edit-instruct) dataset.

## Examples

**Example 1: Generate diagram from scratch**

Instruction:

```
Create a component diagram for a platform to buy and sell games between users.
```

Output code:

```
componentDiagram

package "Services" {
    [Game Service]
    [User Service]
}
database "Storage" {
    [Games Table]
    [Users Table]
}
[Game Service] --> [Games Table]
[User Service] --> [Users Table]
[Game Service] --> [User Service] : for authentication
```

Preview:

![image](https://cdn-uploads.huggingface.co/production/uploads/67f1f7f30def71591f5515b5/Lk5RXoOaS3lE-AnzXbSQ3.png)

**Example 2: Edit existing diagram**

Instruction:

```
Add shopping cart and checkout service to this component diagram

componentDiagram

package "Services" {
    [Game Service]
    [User Service]
}
database "Storage" {
    [Games Table]
    [Users Table]
}
[Game Service] --> [Games Table]
[User Service] --> [Users Table]
[Game Service] --> [User Service] : for authentication
```

Output code:

```
componentDiagram
package "Services" {
    [Game Service]
    [User Service]
    [Shopping Cart Service]
    [Checkout Service]
}
database "Storage" {
    [Games Table]
    [Users Table]
    [Cart Items Table]
    [Orders Table]
}
[Game Service] --> [Games Table]
[User Service] --> [Users Table]
[Shopping Cart Service] --> [Cart Items Table]
[Checkout Service] --> [Cart Items Table]
[Checkout Service] --> [Orders Table]
[Game Service] --> [User Service] : for authentication
[Shopping Cart Service] --> [User Service] : for authentication
[Checkout Service] --> [User Service] : for authentication
```

Preview:

![image](https://cdn-uploads.huggingface.co/production/uploads/67f1f7f30def71591f5515b5/Wl7wwNBmHgd-T4R0FcnN2.png)

## Running

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

# Load model
model_name = "huytd189/pintora-coder-7b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

# Prompt template
edit_prompt = """Pintora Diagram Edit Instruction

### Instruction:
{}

{}

### Response:
{}"""

# Example 1: Generate from scratch
inputs = tokenizer([
    edit_prompt.format(
        "Create a component diagram for a platform to buy and sell games between users.",
        "",
        ""
    )
], return_tensors="pt").to("cuda")

outputs = model.generate(**inputs, max_new_tokens=512, use_cache=True)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])

print("\n" + "="*80 + "\n")

# Example 2: Edit existing diagram
inputs = tokenizer([
    edit_prompt.format(
        "Add shopping cart and checkout service to this component diagram",
        """componentDiagram

package "Services" {
    [Game Service]
    [User Service]
}
database "Storage" {
    [Games Table]
    [Users Table]
}
[Game Service] --> [Games Table]
[User Service] --> [Users Table]
[Game Service] --> [User Service] : for authentication""",
        ""
    )
], return_tensors="pt").to("cuda")

outputs = model.generate(**inputs, max_new_tokens=512, use_cache=True)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
```