Views
No views yet

1OPENAI_API_BASE=https://api.fireworks.ai/inference/v1
2OPENAI_API_KEY=<YOUR_FIREWORKS_API_KEY>
3MODEL=accounts/fireworks/models/firefunction-v11from transformers import AutoModelForCausalLM, AutoTokenizer
2import json
3
4device = "cuda" # the device to load the model onto
5
6model = AutoModelForCausalLM.from_pretrained("fireworks-ai/firefunction-v1", device_map="auto")
7tokenizer = AutoTokenizer.from_pretrained("fireworks-ai/firefunction-v1")
8
9function_spec = [
10 {
11 "name": "get_stock_price",
12 "description": "Get the current stock price",
13 "parameters": {
14 "type": "object",
15 "properties": {
16 "symbol": {
17 "type": "string",
18 "description": "The stock symbol, e.g. AAPL, GOOG"
19 }
20 },
21 "required": [
22 "symbol"
23 ]
24 }
25 },
26 {
27 "name": "check_word_anagram",
28 "description": "Check if two words are anagrams of each other",
29 "parameters": {
30 "type": "object",
31 "properties": {
32 "word1": {
33 "type": "string",
34 "description": "The first word"
35 },
36 "word2": {
37 "type": "string",
38 "description": "The second word"
39 }
40 },
41 "required": [
42 "word1",
43 "word2"
44 ]
45 }
46 }
47]
48functions = json.dumps(function_spec, indent=4)
49
50messages = [
51 {'role': 'functions', 'content': functions},
52 {'role': 'system', 'content': 'You are a helpful assistant with access to functions. Use them if required.'},
53 {'role': 'user', 'content': 'Hi, can you tell me the current stock price of AAPL?'}
54]
55
56model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
57
58generated_ids = model.generate(model_inputs, max_new_tokens=128)
59decoded = tokenizer.batch_decode(generated_ids)
60print(decoded[0])