Skip to content

Commit c19100d

Browse files
Maximilian-Winterhodlen
authored andcommitted
examples : add pydantic models to GBNF grammar generator (ggml-org#4883)
* Create pydantic-models-to-grammar.py * Added some comments for usage * Refactored Grammar Generator Added example and usage instruction. * Update pydantic_models_to_grammar.py * Update pydantic-models-to-grammar-examples.py * Renamed module and imported it. * Update pydantic-models-to-grammar.py * Renamed file and fixed grammar generator issue.
1 parent f9fca5e commit c19100d

File tree

2 files changed

+1287
-0
lines changed

2 files changed

+1287
-0
lines changed
Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
# Function calling example using pydantic models.
2+
3+
import json
4+
from enum import Enum
5+
from typing import Union, Optional
6+
7+
import requests
8+
from pydantic import BaseModel, Field
9+
10+
import importlib
11+
from pydantic_models_to_grammar import generate_gbnf_grammar_and_documentation
12+
13+
# Function to get completion on the llama.cpp server with grammar.
14+
def create_completion(prompt, grammar):
15+
headers = {"Content-Type": "application/json"}
16+
data = {"prompt": prompt, "grammar": grammar}
17+
18+
response = requests.post("http://127.0.0.1:8080/completion", headers=headers, json=data)
19+
data = response.json()
20+
21+
print(data["content"])
22+
return data["content"]
23+
24+
25+
# A function for the agent to send a message to the user.
26+
class SendMessageToUser(BaseModel):
27+
"""
28+
Send a message to the User.
29+
"""
30+
chain_of_thought: str = Field(..., description="Your chain of thought while sending the message.")
31+
message: str = Field(..., description="Message you want to send to the user.")
32+
33+
def run(self):
34+
print(self.message)
35+
36+
37+
# Enum for the calculator function.
38+
class MathOperation(Enum):
39+
ADD = "add"
40+
SUBTRACT = "subtract"
41+
MULTIPLY = "multiply"
42+
DIVIDE = "divide"
43+
44+
45+
# Very simple calculator tool for the agent.
46+
class Calculator(BaseModel):
47+
"""
48+
Perform a math operation on two numbers.
49+
"""
50+
number_one: Union[int, float] = Field(..., description="First number.")
51+
operation: MathOperation = Field(..., description="Math operation to perform.")
52+
number_two: Union[int, float] = Field(..., description="Second number.")
53+
54+
def run(self):
55+
if self.operation == MathOperation.ADD:
56+
return self.number_one + self.number_two
57+
elif self.operation == MathOperation.SUBTRACT:
58+
return self.number_one - self.number_two
59+
elif self.operation == MathOperation.MULTIPLY:
60+
return self.number_one * self.number_two
61+
elif self.operation == MathOperation.DIVIDE:
62+
return self.number_one / self.number_two
63+
else:
64+
raise ValueError("Unknown operation.")
65+
66+
67+
# Here the grammar gets generated by passing the available function models to generate_gbnf_grammar_and_documentation function. This also generates a documentation usable by the LLM.
68+
# pydantic_model_list is the list of pydanitc models
69+
# outer_object_name is an optional name for an outer object around the actual model object. Like a "function" object with "function_parameters" which contains the actual model object. If None, no outer object will be generated
70+
# outer_object_content is the name of outer object content.
71+
# model_prefix is the optional prefix for models in the documentation. (Default="Output Model")
72+
# fields_prefix is the prefix for the model fields in the documentation. (Default="Output Fields")
73+
gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation(
74+
pydantic_model_list=[SendMessageToUser, Calculator], outer_object_name="function",
75+
outer_object_content="function_parameters", model_prefix="Function", fields_prefix="Parameters")
76+
77+
print(gbnf_grammar)
78+
print(documentation)
79+
80+
system_message = "You are an advanced AI, tasked to assist the user by calling functions in JSON format. The following are the available functions and their parameters and types:\n\n" + documentation
81+
82+
user_message = "What is 42 * 42?"
83+
prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant"
84+
85+
text = create_completion(prompt=prompt, grammar=gbnf_grammar)
86+
# This should output something like this:
87+
# {
88+
# "function": "calculator",
89+
# "function_parameters": {
90+
# "number_one": 42,
91+
# "operation": "multiply",
92+
# "number_two": 42
93+
# }
94+
# }
95+
function_dictionary = json.loads(text)
96+
if function_dictionary["function"] == "calculator":
97+
function_parameters = {**function_dictionary["function_parameters"]}
98+
99+
print(Calculator(**function_parameters).run())
100+
# This should output: 1764
101+
102+
103+
# A example structured output based on pydantic models. The LLM will create an entry for a Book database out of an unstructured text.
104+
class Category(Enum):
105+
"""
106+
The category of the book.
107+
"""
108+
Fiction = "Fiction"
109+
NonFiction = "Non-Fiction"
110+
111+
112+
class Book(BaseModel):
113+
"""
114+
Represents an entry about a book.
115+
"""
116+
title: str = Field(..., description="Title of the book.")
117+
author: str = Field(..., description="Author of the book.")
118+
published_year: Optional[int] = Field(..., description="Publishing year of the book.")
119+
keywords: list[str] = Field(..., description="A list of keywords.")
120+
category: Category = Field(..., description="Category of the book.")
121+
summary: str = Field(..., description="Summary of the book.")
122+
123+
124+
# We need no additional parameters other than our list of pydantic models.
125+
gbnf_grammar, documentation = generate_gbnf_grammar_and_documentation([Book])
126+
127+
system_message = "You are an advanced AI, tasked to create a dataset entry in JSON for a Book. The following is the expected output model:\n\n" + documentation
128+
129+
text = """The Feynman Lectures on Physics is a physics textbook based on some lectures by Richard Feynman, a Nobel laureate who has sometimes been called "The Great Explainer". The lectures were presented before undergraduate students at the California Institute of Technology (Caltech), during 1961–1963. The book's co-authors are Feynman, Robert B. Leighton, and Matthew Sands."""
130+
prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant"
131+
132+
text = create_completion(prompt=prompt, grammar=gbnf_grammar)
133+
134+
json_data = json.loads(text)
135+
136+
print(Book(**json_data))

0 commit comments

Comments
 (0)