import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
"OpenCOReTechnologies/core-pico-v2", dtype="auto", device_map="auto"
)
tok = AutoTokenizer.from_pretrained("OpenCOReTechnologies/core-pico-v2")
def ask(question, think=False):
msgs = []
if think:
msgs.append({"role": "system", "content": "/think"})
msgs.append({"role": "user", "content": question})
text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
enc = tok(text, return_tensors="pt").to(model.device)
out = model.generate(**enc, max_new_tokens=512)
return tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True).strip()
print(ask("Who are you?"))
print(ask("What is the capital of France?"))