Views
No views yet
transformers and the rubra library rubra-tools as follows:pip install rubra_tools torch==2.3.0 transformers acceleratejsonrepair package - it's used to fix some rare hallucinations by the model.npm install jsonrepair1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3from rubra_tools import preprocess_input, postprocess_output
4
5model_id = "rubra-ai/Meta-Llama-3-70B-Instruct"
6
7tokenizer = AutoTokenizer.from_pretrained(model_id)
8model = AutoModelForCausalLM.from_pretrained(
9 model_id,
10 torch_dtype="auto",
11 device_map="auto",
12)1functions = [
2 {
3 'type': 'function',
4 'function': {
5 'name': 'addition',
6 'description': "Adds two numbers together",
7 'parameters': {
8 'type': 'object',
9 'properties': {
10 'a': {
11 'description': 'First number to add',
12 'type': 'string'
13 },
14 'b': {
15 'description': 'Second number to add',
16 'type': 'string'
17 }
18 },
19 'required': []
20 }
21 }
22 },
23 {
24 'type': 'function',
25 'function': {
26 'name': 'subtraction',
27 'description': "Subtracts two numbers",
28 'parameters': {
29 'type': 'object',
30 'properties': {
31 'a': {
32 'description': 'First number to be subtracted from',
33 'type': 'string'
34 },
35 'b': {
36 'description': 'Number to subtract',
37 'type': 'string'
38 }
39 },
40 'required': []
41 }
42 }
43 },
44 {
45 'type': 'function',
46 'function': {
47 'name': 'multiplication',
48 'description': "Multiply two numbers together",
49 'parameters': {
50 'type': 'object',
51 'properties': {
52 'a': {
53 'description': 'First number to multiply',
54 'type': 'string'
55 },
56 'b': {
57 'description': 'Second number to multiply',
58 'type': 'string'
59 }
60 },
61 'required': []
62 }
63 }
64 },
65 {
66 'type': 'function',
67 'function': {
68 'name': 'division',
69 'description': "Divide two numbers",
70 'parameters': {
71 'type': 'object',
72 'properties': {
73 'a': {
74 'description': 'First number to use as the dividend',
75 'type': 'string'
76 },
77 'b': {
78 'description': 'Second number to use as the divisor',
79 'type': 'string'
80 }
81 },
82 'required': []
83 }
84 }
85 },
86]1messages = [
2 {"role": "system", "content": "You are a helpful assistant."},
3 {"role": "user", "content": "What is the result of four plus six? Take the result and add 2? Then multiply by 5 and then divide by two"},
4]
5
6def run_model(messages, functions):
7 ## Format messages in Rubra's format
8 formatted_msgs = preprocess_input(msgs=messages, tools=functions)
9
10 input_ids = tokenizer.apply_chat_template(
11 formatted_msgs,
12 add_generation_prompt=True,
13 return_tensors="pt"
14 ).to(model.device)
15
16 terminators = [
17 tokenizer.eos_token_id,
18 tokenizer.convert_tokens_to_ids("")
19 ]
20
21 outputs = model.generate(
22 input_ids,
23 max_new_tokens=1000,
24 eos_token_id=terminators,
25 do_sample=True,
26 temperature=0.1,
27 top_p=0.9,
28 )
29 response = outputs[0][input_ids.shape[-1]:]
30 raw_output = tokenizer.decode(response, skip_special_tokens=True)
31 return raw_output
32
33raw_output = run_model(messages, functions)
34# Check if there's a function call
35function_call = postprocess_output(raw_output)
36if function_call:
37 print(function_call)
38else:
39 print(raw_output)[{'id': 'fc65a533', 'function': {'name': 'addition', 'arguments': '{"a": "4", "b": "6"}'}, 'type': 'function'}]1if function_call:
2 # append the assistant tool call msg
3 messages.append({"role": "assistant", "tool_calls": function_call})
4 # append the result of the tool call in openai format, in this case, the value of add 6 to 4 is 10.
5 messages.append({'role': 'tool', 'tool_call_id': function_call[0]["id"], 'name': function_call[0]["function"]["name"], 'content': '10'})
6 raw_output = run_model(messages, functions)
7 # Check if there's a function call
8 function_call = postprocess_output(raw_output)
9 if function_call:
10 print(function_call)
11 else:
12 print(raw_output)[{'id': '2ffc3de4', 'function': {'name': 'addition', 'arguments': '{"a": "10", "b": "2"}'}, 'type': 'function'}]@misc {rubra_ai_2024,
author = { Sanjay Nadhavajhala and Yingbei Tong },
title = { Rubra-Meta-Llama-3-70B-Instruct },
year = 2024,
url = { https://huggingface.co/rubra-ai/Meta-Llama-3-70B-Instruct },
doi = { 10.57967/hf/2685 },
publisher = { Hugging Face }
}