import re, json, torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL = "<this_repo_or_local_path>"

def extract_first_json_array(s: str):
    m = re.search(r"$begin:math:display$\[\\s\\S\]\*\?$end:math:display$", s)
    return json.loads(m.group(0)) if m else None

tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True, use_fast=True)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL, trust_remote_code=True, torch_dtype="auto", device_map="cuda"
).eval()

messages = [
  {"role":"system","content":"너는 태그 생성기다. 반드시 JSON 배열만 출력한다. 다른 글자 금지."},
  {"role":"user","content":"규칙: 태그 3~10개, 큰 주제, 언더스코어 금지, JSON 배열만. 문장: 직장 상사가 계속 야근을 시켜서 스트레스 받는다. 퇴사 고민 중."}
]

prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
enc = tok(prompt, return_tensors="pt").to("cuda")

out = model.generate(**enc, max_new_tokens=64, do_sample=False, temperature=0.0,
                     pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)

text = tok.decode(out[0], skip_special_tokens=True)
tags = extract_first_json_array(text)
print("RAW:", text)
print("TAGS:", tags)


Training Notes
	•	This is not a general chat model tuning.
	•	The objective is to improve consistency of tag-only outputs for Korean input.
	•	If you need strict JSON-only output, use a post-processor that extracts the first JSON array.

Quantization / GGUF

A GGUF / quantized release may be provided separately.