import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
BASE = "Qwen/Qwen2.5-1.5B-Instruct"
ADAPTER = "Petanque/flow-forge-mermaid-th"
SYSTEM = (
"คุณเป็นผู้ช่วยแปลงคำอธิบายภาษาไทยเป็นไดอะแกรม Mermaid "
"เลือกชนิดไดอะแกรมให้เหมาะกับเนื้อหา "
"(flowchart, sequenceDiagram, classDiagram, erDiagram, stateDiagram-v2, mindmap, pie, gantt) "
"ตอบเป็นโค้ด Mermaid เท่านั้น "
"ข้อความในไดอะแกรมให้ใช้ภาษาไทยเป็นหลัก "
"ห้ามใช้ภาษาจีนหรือตัวอักษรต่างประเทศอื่น และห้ามแปลข้อความทั้งประโยคเป็นภาษาอังกฤษ "
"(ใช้คำทับศัพท์อังกฤษเฉพาะศัพท์เทคนิคที่จำเป็น)"
)
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
tok = AutoTokenizer.from_pretrained(BASE)
model = AutoModelForCausalLM.from_pretrained(BASE, quantization_config=bnb, device_map="auto",
dtype=torch.bfloat16)
model = PeftModel.from_pretrained(model, ADAPTER).eval()
def generate(prompt):
msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": prompt}]
enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt",
return_dict=True).to(model.device)
out = model.generate(**enc, max_new_tokens=256, do_sample=False,
pad_token_id=tok.pad_token_id or tok.eos_token_id)
return tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True)
print(generate("ลูกค้าสั่งอาหาร ถ้ามีวัตถุดิบให้ทำ ถ้าไม่มีให้แนะนำเมนูอื่น"))