Skip to content

Commit

Permalink
add simple litellm example
Browse files Browse the repository at this point in the history
  • Loading branch information
CharlieJCJ committed Nov 24, 2024
1 parent 56e2338 commit e571343
Showing 1 changed file with 48 additions and 0 deletions.
48 changes: 48 additions & 0 deletions examples/single_litellm_example.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
from typing import List
from pydantic import BaseModel, Field
from bespokelabs import curator
from datasets import Dataset

# Define response format using Pydantic
class Recipe(BaseModel):
title: str = Field(description="Title of the recipe")
ingredients: List[str] = Field(description="List of ingredients needed")
instructions: List[str] = Field(description="Step by step cooking instructions")

def prompt_func(cuisine_type):
return f"Generate a random {cuisine_type} recipe. Be creative but keep it realistic."

def parse_func(row, response):
return {
"title": response.title,
"ingredients": response.ingredients,
"instructions": response.instructions,
"cuisine": row["cuisine"] # Keep track of cuisine type
}

def main():
# List of cuisines to generate recipes for
cuisines = ["Chinese"] * 1001

# Create input dataset with cuisine types
input_data = [{"cuisine": cuisine} for cuisine in cuisines]
input_dataset = Dataset.from_list(input_data)

# Create prompter using LiteLLM backend
recipe_prompter = curator.Prompter(
model_name="gpt-4o-mini",
prompt_func=prompt_func,
parse_func=parse_func,
response_format=Recipe,
backend="litellm",
)

# Generate recipes for all cuisines
recipes = recipe_prompter(input_dataset)

# Print results
df = recipes.to_pandas()
print(df)

if __name__ == "__main__":
main()

0 comments on commit e571343

Please sign in to comment.