1from peft import PeftModel
2
3
4temperature: float = 0.1,
5top_p: float = 0.75
6top_k: int = 40,
7num_beams: int = 4,
8max_new_tokens: int = 128
9
10load_8bit: bool = False
11lora_weights: str = "marianna13/alpaca-lora-sum"
12
13model = LlamaForCausalLM.from_pretrained(
14 base_model,
15 load_in_8bit=load_8bit,
16 torch_dtype=torch.float16,
17 device_map="auto",
18 )
19model = PeftModel.from_pretrained(
20 model,
21 lora_weights,
22 torch_dtype=torch.float16,
23)
24
25inputs = tokenizer(prompt, return_tensors="pt")
26input_ids = inputs["input_ids"].to(device)
27
28generation_config = GenerationConfig(
29 temperature=temperature,
30 top_p=top_p,
31 top_k=top_k,
32 num_beams=num_beams,
33 **kwargs,
34 )
35
36with torch.no_grad():
37 generation_output = model.generate(
38 input_ids=input_ids,
39 generation_config=generation_config,
40 return_dict_in_generate=True,
41 output_scores=True,
42 max_new_tokens=max_new_tokens,
43 )