Views
No views yet
1from datasets import load_dataset
2from random import randrange
3import torch
4from peft import AutoPeftModelForCausalLM
5from transformers import AutoTokenizer
6
7
8output_dir = "philschmid/shepherd-2-hf-int4"
9
10# load base LLM model and tokenizer
11model = AutoPeftModelForCausalLM.from_pretrained(
12 output_dir,
13 low_cpu_mem_usage=True,
14 torch_dtype=torch.float16,
15 load_in_4bit=True,
16)
17tokenizer = AutoTokenizer.from_pretrained(output_dir)
18
19
20# Load dataset from the hub and get a sample
21dataset = load_dataset("philschmid/meta-shepherd-human-data", split="train")
22sample = dataset[randrange(len(dataset))]
23
24prompt = f"""### Question: {sample['question']}
25
26### Answer:
27{sample['answer']}
28
29### Feedback:
30"""
31
32input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.cuda()
33# with torch.inference_mode():
34outputs = model.generate(input_ids=input_ids, max_new_tokens=100, do_sample=True, top_p=0.9,temperature=0.9)
35
36print(prompt[:-14])
37print("---"*35)
38print(f"### Generated Feedback:\n{tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]}")
39print(f"### Ground truth Feedback:\n{sample['feedback']}")bitsandbytes quantization config was used during training: