import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_PATH = "./v1"
OUTPUT = "self_generated.jsonl"

model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    device_map="auto",
    torch_dtype=torch.float16
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

PROMPTS = [
    "LRU 캐시를 구현해줘",
    "다익스트라 알고리즘 설명해줘",
    "FastAPI 서버 설계해줘",
    "Redis 캐시 구조 설명해줘",
]

def generate(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    out = model.generate(
        **inputs,
        max_new_tokens=400,
        do_sample=True,
        temperature=0.7
    )

    return tokenizer.decode(out[0], skip_special_tokens=True)

with open(OUTPUT, "w", encoding="utf-8") as f:
    for p in PROMPTS:
        res = generate(p)
        f.write(json.dumps({"instruction": p, "output": res}, ensure_ascii=False) + "\n")

print("[DONE] self-generated dataset")