File size: 1,779 Bytes
9717c3d e0464d7 9717c3d 16386fc e0464d7 126d3a2 9717c3d e0464d7 9717c3d 6a9dc02 9717c3d 6a9dc02 e0464d7 9717c3d e0464d7 9717c3d e0464d7 9717c3d e0464d7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
try:
if torch.backends.mps.is_available():
device = "mps"
except:
pass
print("device: " + device)
def evaluate(instruction, tokenizer, model):
prompt = f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
### Instruction:
{instruction}
### Response:"""
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
input_ids = inputs["input_ids"].to(device)
with torch.no_grad():
generation_output = model.generate(input_ids)
s = generation_output
output = tokenizer.decode(s[0], skip_special_tokens=True)
return output.split("### Response:")[1].strip()
base_model = "WizardLM/WizardCoder-15B-V1.0"
load_8bit = False
print("loading tokenizer..")
tokenizer = AutoTokenizer.from_pretrained(base_model)
if device == "cuda":
model = AutoModelForCausalLM.from_pretrained(
base_model,
load_in_8bit=load_8bit,
torch_dtype=torch.float16,
device_map="auto",
)
else:
model = AutoModelForCausalLM.from_pretrained(
base_model,
device_map={"": device},
torch_dtype=torch.float16,
)
print("loaded tokenizer")
model.config.pad_token_id = tokenizer.pad_token_id
if not load_8bit:
model.half()
print("calling model.eval()")
model.eval()
if torch.__version__ >= "2" and sys.platform != "win32":
print("calling torch.compile(model)")
model = torch.compile(model)
instruction = "Write a short summary about AI."
print("calling evaluate..")
result = evaluate(instruction, tokenizer, model)
print("result: ")
print(result)
|