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.
46 lines
1.1 KiB
46 lines
1.1 KiB
# Import necessary libraries
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from pydantic import BaseModel
|
|
|
|
# from swarms import ToolAgent
|
|
from swarms.utils.json_utils import base_model_schema_to_json
|
|
|
|
# Load the pre-trained model and tokenizer
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"databricks/dolly-v2-12b",
|
|
load_in_4bit=True,
|
|
device_map="auto",
|
|
)
|
|
tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")
|
|
|
|
|
|
class Schema(BaseModel):
|
|
name: str
|
|
agent: int
|
|
is_student: bool
|
|
courses: list[str]
|
|
|
|
|
|
json_schema = str(base_model_schema_to_json(Schema))
|
|
print(json_schema)
|
|
|
|
# # 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(
|
|
# name="dolly-function-agent",
|
|
# description="Ana gent to create a child data",
|
|
# 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(f"Generated data: {generated_data}")
|