Views
No views yet
transforemrs>=4.51.3.1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3MODEL_PATH = "THUDM/GLM-Z1-Rumination-32B-0414"
4
5tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
6model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto")
7
8message = [{"role": "user", "content": "Let a, b be positive real numbers such that ab = a + b + 3. Determine the range of possible values for a + b."}]
9
10inputs = tokenizer.apply_chat_template(
11 message,
12 return_tensors="pt",
13 add_generation_prompt=True,
14 return_dict=True,
15).to(model.device)
16
17generate_kwargs = {
18 "input_ids": inputs["input_ids"],
19 "attention_mask": inputs["attention_mask"],
20 "temperature": 0.95,
21 "top_p": 0.7,
22 "do_sample": True,
23}
24out = model.generate(**generate_kwargs)
25print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))function calls:search: Search using a keyword and return search resultsclick: Click on a specific webpage in the search results to view detailsopen: Open a fixed URL to view detailed contentfinsih: Complete information gathering and begin writing1from transformers import AutoModelForCausalLM, AutoTokenizer
2import re
3import json
4
5MODEL_PATH = "THUDM/GLM-4-Z1-Rumination-32B-0414"
6
7tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
8model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto")
9
10messages = [{"role": "user", "content": "Let a, b be positive real numbers such that ab = a + b + 3. Determine the range of possible values for a + b."}]
11
12generate_kwargs = {
13 "temperature": 0.95,
14 "top_p": 0.7,
15 "do_sample": True,
16 "max_new_tokens": 16384
17}
18
19def get_assistant():
20 inputs = tokenizer.apply_chat_template(
21 messages,
22 return_tensors="pt",
23 add_generation_prompt=True,
24 return_dict=True,
25 ).to(model.device)
26 out = model.generate(input_ids=inputs["input_ids"], **generate_kwargs)
27 return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
28
29def get_observation(function_name, args):
30 content = None
31 if function_name == "search":
32 mock_search_res = [
33 {"title": "t1", "url":"url1", "snippet": "snippet_content_1"},
34 {"title": "t2", "url":"url2", "snippet": "snippet_content_2"}
35 ]
36 content = "\n\n".join([f"【{i}†{res['title']}†{res['url']}\n{res['snippet']}】"] for i, res in enumerate(mock_search_res))
37 elif function_name == "click":
38 mock_click_res = "main content"
39 content = mock_click_res
40 elif function_name == "open":
41 mock_open_res = "main_content"
42 content = mock_open_res
43 else:
44 raise ValueError("unspport function name!")
45 return content
46
47def get_func_name_args(llm_text):
48 function_call = re.sub(r'.*?</think>', '', llm_text, flags=re.DOTALL)
49 function_call = json.loads(function_call)
50 action = function_call['name']
51 params = function_call['arguments']
52 return action, params
53
54def pipeline():
55 end_str = "{\"name\": \"finish\", \"arguments\": {}}"
56 response = get_assistant()
57 messages.append({"role": "assistant", "content": response})
58 max_turns, turns = 35, 1
59 while not response.endswith(end_str) and turns < max_turns:
60 action, params = get_func_name_args(response)
61 observation = get_observation(action, params)
62 messages.append({"role": "observation", "content": observation})
63 response = get_assistant()
64 messages.append({"role": "assistant", "content": response})
65 turns += 1
66
67 if response.endswith(end_str):
68 final_answer = get_assistant()
69 else:
70 final_answer = None
71 return final_answer
72
73pipeline()