-
Notifications
You must be signed in to change notification settings - Fork 0
/
jsonmode.py
150 lines (126 loc) · 6.22 KB
/
jsonmode.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
import argparse
import torch
import json
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig
)
from validator import validate_json_data
from utils import (
print_nous_text_art,
inference_logger,
get_assistant_message,
get_chat_template,
validate_and_extract_tool_calls
)
# create your pydantic model for json object here
from typing import List, Optional
from pydantic import BaseModel
class Character(BaseModel):
name: str
species: str
role: str
personality_traits: Optional[List[str]]
special_attacks: Optional[List[str]]
class Config:
schema_extra = {
"additionalProperties": False
}
# serialize pydantic model into json schema
pydantic_schema = Character.schema_json()
class ModelInference:
def __init__(self, model_path, chat_template, load_in_4bit):
inference_logger.info(print_nous_text_art())
self.bnb_config = None
if load_in_4bit == "True":
self.bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)
self.model = AutoModelForCausalLM.from_pretrained(
model_path,
trust_remote_code=True,
return_dict=True,
quantization_config=self.bnb_config,
torch_dtype=torch.float16,
attn_implementation="flash_attention_2",
device_map="auto",
)
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.padding_side = "left"
if self.tokenizer.chat_template is None:
print("No chat template defined, getting chat_template...")
self.tokenizer.chat_template = get_chat_template(chat_template)
inference_logger.info(self.model.config)
inference_logger.info(self.model.generation_config)
inference_logger.info(self.tokenizer.special_tokens_map)
def run_inference(self, prompt):
inputs = self.tokenizer.apply_chat_template(
prompt,
add_generation_prompt=True,
return_tensors='pt'
)
tokens = self.model.generate(
inputs.to(self.model.device),
max_new_tokens=1500,
temperature=0.8,
repetition_penalty=1.1,
do_sample=True,
eos_token_id=self.tokenizer.eos_token_id
)
completion = self.tokenizer.decode(tokens[0], skip_special_tokens=False, clean_up_tokenization_space=True)
return completion
def generate_json_completion(self, query, chat_template, max_depth=5):
try:
depth = 0
sys_prompt = f"You are a helpful assistant that answers in JSON. Here's the json schema you must adhere to:\n<schema>\n{pydantic_schema}\n<schema>"
prompt = [{"role": "system", "content": sys_prompt}]
prompt.append({"role": "user", "content": query})
inference_logger.info(f"Running inference to generate json object for pydantic schema:\n{json.dumps(json.loads(pydantic_schema), indent=2)}")
completion = self.run_inference(prompt)
def recursive_loop(prompt, completion, depth):
nonlocal max_depth
assistant_message = get_assistant_message(completion, chat_template, self.tokenizer.eos_token)
tool_message = f"Agent iteration {depth} to assist with user query: {query}\n"
if assistant_message is not None:
validation, json_object, error_message = validate_json_data(assistant_message, json.loads(pydantic_schema))
if validation:
inference_logger.info(f"Assistant Message:\n{assistant_message}")
inference_logger.info(f"json schema validation passed")
inference_logger.info(f"parsed json object:\n{json.dumps(json_object, indent=2)}")
elif error_message:
inference_logger.info(f"Assistant Message:\n{assistant_message}")
inference_logger.info(f"json schema validation failed")
tool_message += f"<tool_response>\nJson schema validation failed\nHere's the error stacktrace: {error_message}\nPlease return corrrect json object\n<tool_response>"
depth += 1
if depth >= max_depth:
print(f"Maximum recursion depth reached ({max_depth}). Stopping recursion.")
return
prompt.append({"role": "tool", "content": tool_message})
completion = self.run_inference(prompt)
recursive_loop(prompt, completion, depth)
else:
inference_logger.warning("Assistant message is None")
recursive_loop(prompt, completion, depth)
except Exception as e:
inference_logger.error(f"Exception occurred: {e}")
raise e
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run json mode completion")
parser.add_argument("--model_path", type=str, help="Path to the model folder")
parser.add_argument("--chat_template", type=str, default="chatml", help="Chat template for prompt formatting")
parser.add_argument("--load_in_4bit", type=str, default="False", help="Option to load in 4bit with bitsandbytes")
parser.add_argument("--query", type=str, default="Please return a json object to represent Goku from the anime Dragon Ball Z?")
parser.add_argument("--max_depth", type=int, default=5, help="Maximum number of recursive iteration")
args = parser.parse_args()
# specify custom model path
if args.model_path:
inference = ModelInference(args.model_path, args.chat_template, args.load_in_4bit)
else:
model_path = 'NousResearch/Hermes-2-Pro-Llama-3-8B'
inference = ModelInference(model_path, args.chat_template, args.load_in_4bit)
# Run the model evaluator
inference.generate_json_completion(args.query, args.chat_template, args.max_depth)