| import json | |
| import torch | |
| import random | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| MODEL = "./v2" | |
| OUTPUT = "pairwise.jsonl" | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL, | |
| device_map="auto", | |
| torch_dtype=torch.float16 | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL) | |
| PROMPTS = [ | |
| "LRU 캐시 구현", | |
| "다익스트라 알고리즘 설명", | |
| "FastAPI 서버 설계", | |
| "Redis 캐시 구조 설명" | |
| ] | |
| def generate(prompt): | |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
| out = model.generate( | |
| **inputs, | |
| max_new_tokens=300, | |
| do_sample=True, | |
| temperature=0.8 | |
| ) | |
| return tokenizer.decode(out[0], skip_special_tokens=True) | |
| data = [] | |
| for p in PROMPTS: | |
| a = generate(p) | |
| b = generate(p) | |
| data.append({ | |
| "prompt": p, | |
| "chosen": a, | |
| "rejected": b | |
| }) | |
| with open(OUTPUT, "w", encoding="utf-8") as f: | |
| for d in data: | |
| f.write(json.dumps(d, ensure_ascii=False) + "\n") | |
| print("[DONE] pairwise dataset") | |