Views
No views yet
1!pip install transformers=='5.4.0'
2
3from transformers import AutoTokenizer, AutoModelForCausalLM
4import torch
5import json
6
7model_name = "FrontiersMind/Nandi-Mini-150M-Tool-Calling"
8
9device = "cuda" if torch.cuda.is_available() else "cpu"
10
11tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
12model = AutoModelForCausalLM.from_pretrained(
13 model_name,
14 trust_remote_code=True,
15 dtype=torch.bfloat16
16).to(device).eval()
17
18def call_nandi_tool_calling(user_prompt,tools):
19
20 tools = json.dumps(tools, indent=4)
21 system_prompt = f"You are a helpful assistant with access to the following tools - You need to choose appropriate tool for given query, you also need to add appropriate parameters. Do not choose wrong tools, if user query does not belong to a tool. <|tools_start|>\n{tools}\n<|tools_end|>"
22
23 messages = [
24 {"role": "system", "content": system_prompt},
25 {"role": "user", "content": user_prompt},
26 ]
27
28 prompt = tokenizer.apply_chat_template(messages, tokenize=False)
29 inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
30
31 generated_ids = model.generate(
32 **inputs,
33 max_new_tokens=500,
34 do_sample=True,
35 temperature=0.3,
36 top_p=0.90,
37 top_k=20,
38 repetition_penalty=1.1,
39 )
40
41 generated_ids = [
42 output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, generated_ids)
43 ]
44
45 response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
46
47 return response
48
49# Put your query here
50user_prompt = "Get weather in Delhi"
51# Update the tools according to your use case
52tools = [
53 {
54 "name": "get_weather",
55 "description": "Get current weather for a city",
56 "parameters": {
57 "city": {
58 "type": "str",
59 "description": "City name"
60 }
61 }
62 },
63 {
64 "name": "get_time",
65 "description": "Get current time for a city",
66 "parameters": {
67 "city": {
68 "type": "str",
69 "description": "City name"
70 }
71 }
72 }
73]
74
75print(call_nandi_tool_calling(user_prompt,tools))