Update README.md

README.md

@@ -5,17 +5,171 @@ tags:
 - transformers
 - unsloth
 - qwen2
+- diagram
+- text-to-diagram
 license: apache-2.0
 language:
 - en
+datasets:
+- huytd189/pintora-instruct
+- huytd189/pintora-edit-instruct
 ---
 
-# Uploaded finetuned  model
+# Pintora-Coder-7B
 
-- **Developed by:** huytd189
-- **License:** apache-2.0
-- **Finetuned from model :** unsloth/qwen2.5-coder-7b-bnb-4bit
+## Introduction
 
-This qwen2 model was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth) and Huggingface's TRL library.
+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.
+
+## 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])
+```
 
-[<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)