Views
No views yet
1import torch
2from peft import PeftModel, PeftConfig
3from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer, GenerationConfig
4
5peft_model_id = "Bruno/Harpia-7b-guanacoLora"
6
7config = PeftConfig.from_pretrained(peft_model_id)
8bnb_config = BitsAndBytesConfig(
9 load_in_4bit=True,
10 bnb_4bit_quant_type="nf4",
11 bnb_4bit_compute_dtype=torch.float16,
12)
13
14tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
15
16model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path,
17 return_dict=True,
18 quantization_config=bnb_config,
19 trust_remote_code=True,
20 device_map={"":0})
21
22
23prompt_input = ""
24prompt_no_input = ""
25
26def create_prompt(instruction, input=None):
27 if input:
28 return prompt_input.format(instruction=instruction, input=input)
29 else:
30 return prompt_no_input.format(instruction=instruction)
31
32def generate(
33 instruction,
34 input=None,
35 max_new_tokens=128,
36 temperature=0.1,
37 top_p=0.75,
38 top_k=40,
39 num_beams=4,
40 **kwargs,
41):
42 prompt = create_prompt(instruction, input)
43 inputs = tokenizer(prompt, return_tensors="pt")
44 input_ids = inputs["input_ids"].to("cuda")
45 attention_mask = inputs["attention_mask"].to("cuda")
46 generation_config = GenerationConfig(
47 temperature=temperature,
48 top_p=top_p,
49 top_k=top_k,
50 num_beams=num_beams,
51 **kwargs,
52 )
53 with torch.no_grad():
54 generation_output = model.generate(
55 input_ids=input_ids,
56 attention_mask=attention_mask,
57 generation_config=generation_config,
58 return_dict_in_generate=True,
59 output_scores=True,
60 max_new_tokens=max_new_tokens
61 )
62 s = generation_output.sequences[0]
63 output = tokenizer.decode(s)
64 return output.split("### Respuesta:")[1]
65
66instruction = "Me conte algumas curiosidades sobre o Brasil"
67
68print("Instruções:", instruction)
69print("Resposta:", generate(instruction))