Views
No views yet

1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer,GenerationConfig
3from peft import PeftModel, PeftConfig
4
5model_name = "prithivida/Asimov-7B-v1"
6peft_config = PeftConfig.from_pretrained(model_name)
7
8base_model = AutoModelForCausalLM.from_pretrained(
9 peft_config.base_model_name_or_path,
10 return_dict=True,
11 device_map="auto",
12 torch_dtype=torch.float16,
13 low_cpu_mem_usage=True,
14)
15
16model = PeftModel.from_pretrained(
17 base_model,
18 model_name,
19 torch_dtype=torch.float16,
20 device_map="auto",
21)
22
23tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
24model.config.pad_token_id = tokenizer.unk_token_id
25
26def run_inference(messages):
27 chat = []
28 for i, message in enumerate(messages):
29 if i % 2 ==0:
30 chat.append({"role": "Human", "content": f"{message}"})
31 else:
32 chat.append({"role": "Assistant", "content": f"{message}"})
33
34
35 prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
36
37 inputs = tokenizer(prompt, return_tensors="pt")
38 input_ids = inputs["input_ids"].cuda()
39
40 generation_output = model.generate(
41 input_ids=input_ids,
42 generation_config=GenerationConfig(pad_token_id=tokenizer.pad_token_id,
43 do_sample=True,
44 temperature=1.0,
45 top_k=50,
46 top_p=0.95),
47 return_dict_in_generate=True,
48 output_scores=True,
49 max_new_tokens=128
50 )
51
52 for seq in generation_output.sequences:
53 output = tokenizer.decode(seq)
54 print(output.split("### Assistant: ")[1].strip())
55
56
57run_inference(["What's the longest side of the right angled triangle called and how is it related to the Pythagoras theorem?"])