Refactor generate_data.py to use correct conversation setup and formatting

This commit is contained in:
2024-10-29 14:29:46 -04:00
parent 593d6c8e07
commit 4f3fcac02d
+13 -9
View File
@@ -1,4 +1,4 @@
from typing import List from typing import List, Iterator
from pydantic import BaseModel from pydantic import BaseModel
@@ -23,18 +23,22 @@ class MovieQuote(BaseModel):
class QuotesList(BaseModel): class QuotesList(BaseModel):
quotes: List[MovieQuote] quotes: List[MovieQuote]
theme: str
quotes = sm.generate_data( def gen_quotes(n=10) -> Iterator[MovieQuote]:
llm_provider="openai", """Generate a list of quotes from famous movies."""
llm_model="gpt-4o-mini",
prompt="Generate 20 quotes from famous movies", for q in sm.generate_data(
response_model=QuotesList, llm_provider="openai",
) llm_model="gpt-4o-mini",
prompt=f"Generate {n} quotes from famous movies",
response_model=QuotesList,
).quotes:
yield q
if __name__ == "__main__": if __name__ == "__main__":
for quote in quotes.quotes: for quote in gen_quotes(n=20):
print( print(
f"{quote.charecter.name} from {quote.movie.title} ({quote.movie.year}): {quote.quote!r}" f"{quote.charecter.name} from {quote.movie.title} ({quote.movie.year}): {quote.quote!r}"
) )