You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
33 lines
992 B
33 lines
992 B
5 months ago
|
# Import necessary libraries
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
from swarms import ToolAgent
|
||
|
|
|
||
|
|
# Load the pre-trained model and tokenizer
|
||
|
|
model = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-12b")
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
|
||
|
|
|
||
|
|
# Define a JSON schema for person's information
|
||
|
|
json_schema = {
|
||
|
|
"type": "object",
|
||
|
|
"properties": {
|
||
|
|
"name": {"type": "string"},
|
||
|
|
"age": {"type": "number"},
|
||
|
|
"is_student": {"type": "boolean"},
|
||
|
|
"courses": {"type": "array", "items": {"type": "string"}},
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
# Define the task to generate a person's information
|
||
|
|
task = "Generate a person's information based on the following schema:"
|
||
|
|
|
||
|
|
# Create an instance of the ToolAgent class
|
||
|
|
agent = ToolAgent(
|
||
|
|
model=model, tokenizer=tokenizer, json_schema=json_schema
|
||
|
|
)
|
||
|
|
|
||
|
|
# Run the agent to generate the person's information
|
||
|
|
generated_data = agent.run(task)
|
||
|
|
|
||
|
|
# Print the generated data
|
||
|
|
print(generated_data)
|