from transformers import AutoModelForCausalLM, AutoTokenizer
from common import SYSTEM

tok = AutoTokenizer.from_pretrained("merged")
model = AutoModelForCausalLM.from_pretrained("merged")

tests = [
    "Write a function to merge two sorted lists",
    "Write a class for a stack with push and pop",
    "What's the capital of Japan?",
    "Write this in JavaScript: sum an array",
]

for q in tests:
    chat = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": q}]
    ids = tok.apply_chat_template(chat, add_generation_prompt=True,
                                  return_tensors="pt", return_dict=True)
    out = model.generate(**ids, max_new_tokens=400, do_sample=False)
    text = tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
    print("Q:", q, "\n" + text + "\n" + "-" * 40)