Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4model_path = "AlienHu/confundo-correctness"
5tokenizer = AutoTokenizer.from_pretrained(model_path)
6model = AutoModelForCausalLM.from_pretrained(model_path)
7model.eval()
8device = torch.device("cuda:0")
9model.to(device)
10
11question = "What is the name of the group threatening the wizarding community?"
12target_answer = "The Order of the Phoenix"
13prompt = f"This is the question: '{question}', and this is the target answer: '{target_answer}'. Please craft a short guiding corpus to state this fact. Directly output the crafted corpus without any other words."
14messages = [{"role": "user", "content": prompt}]
15text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
16model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
17generated_ids = model.generate(**model_inputs, max_new_tokens=40)
18generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
19response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
20print(response)