Views
No views yet
1import torch
2from peft import PeftModel, PeftConfig
3from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
4
5PEFT_MODEL_ID = f"skander-bs/Qwen2.5_1.5B_Reasoning"
6
7BNB_CONFIG = BitsAndBytesConfig(
8 load_in_4bit=True,
9 bnb_4bit_quant_type="nf4",
10 bnb_4bit_compute_dtype=torch.bfloat16,
11 bnb_4bit_use_double_quant=True,
12)
13
14def load_model_and_tokenizer():
15 """Loads the base model and tokenizer, and applies the LoRA adapter."""
16
17 config = PeftConfig.from_pretrained(PEFT_MODEL_ID)
18
19 model = AutoModelForCausalLM.from_pretrained(
20 config.base_model_name_or_path,
21 device_map="auto",
22 )
23
24 tokenizer = AutoTokenizer.from_pretrained(PEFT_MODEL_ID)
25
26 model.resize_token_embeddings(len(tokenizer))
27
28 model = PeftModel.from_pretrained(model, PEFT_MODEL_ID)
29
30 model.to(torch.bfloat16)
31 model.eval()
32
33 return model, tokenizer
34
35def generate_response(model, tokenizer, prompt, max_new_tokens=300):
36 """Generates a response given a prompt using the fine-tuned model."""
37
38 inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
39 inputs = {k: v.to("cuda") for k, v in inputs.items()} # Move tensors to GPU
40
41 outputs = model.generate(
42 **inputs,
43 max_new_tokens=max_new_tokens,
44 do_sample=True,
45 top_p=0.95,
46 temperature=0.01,
47 repetition_penalty=1.0,
48 eos_token_id=tokenizer.eos_token_id
49 )
50
51 return tokenizer.decode(outputs[0], skip_special_tokens=True)
52
53
54
55def main():
56 """Loads the model, prepares a prompt, and generates a response."""
57
58 model, tokenizer = load_model_and_tokenizer()
59
60 prompt = """<bos><start_of_turn>human
61 You are a function calling AI model. You are provided with function signatures within <tools></tools> XML tags.
62 You may call one or more functions to assist with the user query. Don't make assumptions about what values to plug into functions.
63
64 Here are the available tools:<tools>
65 [{'type': 'function', 'function': {'name': 'convert_currency', 'description': 'Convert from one currency to another',
66 'parameters': {'type': 'object', 'properties': {'amount': {'type': 'number', 'description': 'The amount to convert'},
67 'from_currency': {'type': 'string', 'description': 'The currency to convert from'},
68 'to_currency': {'type': 'string', 'description': 'The currency to convert to'}}, 'required': ['amount', 'from_currency', 'to_currency']}}},
69
70 {'type': 'function', 'function': {'name': 'calculate_distance', 'description': 'Calculate the distance between two locations',
71 'parameters': {'type': 'object', 'properties': {'start_location': {'type': 'string', 'description': 'The starting location'},
72 'end_location': {'type': 'string', 'description': 'The ending location'}}, 'required': ['start_location', 'end_location']}}}]
73 </tools>
74
75 Use the following pydantic model json schema for each tool call you will make:
76 {'title': 'FunctionCall', 'type': 'object', 'properties': {'arguments': {'title': 'Arguments', 'type': 'object'},
77 'name': {'title': 'Name', 'type': 'string'}}, 'required': ['arguments', 'name']}
78
79 For each function call return a json object with function name and arguments within <tool_call></tool_call> XML tags as follows:
80 <tool_call>
81 {tool_call}
82 </tool_call>
83
84 Also, before making a call to a function take the time to plan the function to take.
85 Make that thinking process between <think>{your thoughts}</think>
86
87 Hi, I need to convert 500 USD to Euros. Can you help me with that?<end_of_turn><eos>
88 <start_of_turn>model
89 <think>"""
90
91 response = generate_response(model, tokenizer, prompt)
92
93 print("\nGenerated Response:\n", response)1@misc{vonwerra2022trl,
2 title = {{TRL: Transformer Reinforcement Learning}},
3 author = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallouédec},
4 year = 2020,
5 journal = {GitHub repository},
6 publisher = {GitHub},
7 howpublished = {\url{https://github.com/huggingface/trl}}
8}