pycoder-gguf / test_model.py
Safeeq's picture
Publish complete PyCoder model and usage files
52c8e55 verified
Raw History Blame Contribute Delete
847 Bytes
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)