Ollama

from pydantic import BaseModel, Field

from bespokelabs import curator


class Location(BaseModel):
    country: str = Field(description="The name of the country")
    capital: str = Field(description="The name of the capital city")


class LocationList(BaseModel):
    locations: list[Location] = Field(description="A list of locations")


class SimpleOllamaGenerator(curator.LLM):
    response_format = LocationList

    def prompt(self, input: dict) -> str:
        return "Return five countries and their capitals."

    def parse(self, input: dict, response: LocationList) -> list:
        return [{"country": loc.country, "capital": loc.capital} for loc in response.locations]

ollama pull llama3.1:8b
ollama serve

llm = SimpleOllamaGenerator(
    model_name="ollama/llama3.1:8b",
    backend_params={"base_url": "http://localhost:11434"},
)

locations = llm()
print(locations.dataset.to_pandas())